Make the estimated size in ARP a little more accurate.

Fixes 4039

17 files changed, 177 insertions, 113 deletions

diff --git a/src/burn/engine/apply.cpp b/src/burn/engine/apply.cpp
index 9ee7b58c..69106f65 100644
--- a/src/burn/engine/apply.cpp
+++ b/src/burn/engine/apply.cpp
@@ -79,7 +79,8 @@ static HRESULT WINAPI AuthenticationRequired(
 static void CalculateKeepRegistration(
     __in BURN_ENGINE_STATE* pEngineState,
     __in BOOL fLog,
-    __inout BOOTSTRAPPER_REGISTRATION_TYPE* pRegistrationType
+    __inout BOOTSTRAPPER_REGISTRATION_TYPE* pRegistrationType,
+    __inout DWORD64* pqwEstimatedSize
     );
 static HRESULT ExecuteDependentRegistrationActions(
     __in HANDLE hPipe,
@@ -422,8 +423,9 @@ extern "C" HRESULT ApplyRegister(
     HRESULT hr = S_OK;
     LPWSTR sczEngineWorkingPath = NULL;
     BOOTSTRAPPER_REGISTRATION_TYPE registrationType = BOOTSTRAPPER_REGISTRATION_TYPE_INPROGRESS;
+    DWORD64 qwEstimatedSize = 0;
 
-    CalculateKeepRegistration(pEngineState, FALSE, &registrationType);
+    CalculateKeepRegistration(pEngineState, FALSE, &registrationType, &qwEstimatedSize);
 
     hr = UserExperienceOnRegisterBegin(&pEngineState->userExperience, &registrationType);
     ExitOnRootFailure(hr, "BA aborted register begin.");
@@ -451,12 +453,12 @@ extern "C" HRESULT ApplyRegister(
         // begin new session
         if (pEngineState->registration.fPerMachine)
         {
-            hr = ElevationSessionBegin(pEngineState->companionConnection.hPipe, sczEngineWorkingPath, pEngineState->registration.sczResumeCommandLine, pEngineState->registration.fDisableResume, &pEngineState->variables, pEngineState->plan.dwRegistrationOperations, pEngineState->registration.fDetectedForeignProviderKeyBundleId, pEngineState->plan.qwEstimatedSize, registrationType);
+            hr = ElevationSessionBegin(pEngineState->companionConnection.hPipe, sczEngineWorkingPath, pEngineState->registration.sczResumeCommandLine, pEngineState->registration.fDisableResume, &pEngineState->variables, pEngineState->plan.dwRegistrationOperations, pEngineState->registration.fDetectedForeignProviderKeyBundleId, qwEstimatedSize, registrationType);
             ExitOnFailure(hr, "Failed to begin registration session in per-machine process.");
         }
         else
         {
-            hr = RegistrationSessionBegin(sczEngineWorkingPath, &pEngineState->registration, &pEngineState->cache, &pEngineState->variables, pEngineState->plan.dwRegistrationOperations, pEngineState->plan.qwEstimatedSize, registrationType);
+            hr = RegistrationSessionBegin(sczEngineWorkingPath, &pEngineState->registration, &pEngineState->cache, &pEngineState->variables, pEngineState->plan.dwRegistrationOperations, qwEstimatedSize, registrationType);
             ExitOnFailure(hr, "Failed to begin registration session.");
         }
     }
@@ -491,6 +493,7 @@ extern "C" HRESULT ApplyUnregister(
     BURN_RESUME_MODE resumeMode = BURN_RESUME_MODE_NONE;
     BOOTSTRAPPER_REGISTRATION_TYPE defaultRegistrationType = BOOTSTRAPPER_REGISTRATION_TYPE_NONE;
     BOOTSTRAPPER_REGISTRATION_TYPE registrationType = BOOTSTRAPPER_REGISTRATION_TYPE_NONE;
+    DWORD64 qwEstimatedSize = 0;
 
     // Calculate special cases for the resume mode. If a restart has been initiated, that trumps all other
     // modes. If the user chose to suspend the install then we'll use that as the resume mode.
@@ -513,7 +516,7 @@ extern "C" HRESULT ApplyUnregister(
         defaultRegistrationType = BOOTSTRAPPER_REGISTRATION_TYPE_INPROGRESS;
     }
 
-    CalculateKeepRegistration(pEngineState, TRUE, &defaultRegistrationType);
+    CalculateKeepRegistration(pEngineState, TRUE, &defaultRegistrationType, &qwEstimatedSize);
 
     registrationType = defaultRegistrationType;
 
@@ -547,12 +550,12 @@ extern "C" HRESULT ApplyUnregister(
 
     if (pEngineState->registration.fPerMachine)
     {
-        hr = ElevationSessionEnd(pEngineState->companionConnection.hPipe, resumeMode, restart, pEngineState->registration.fDetectedForeignProviderKeyBundleId, registrationType);
+        hr = ElevationSessionEnd(pEngineState->companionConnection.hPipe, resumeMode, restart, pEngineState->registration.fDetectedForeignProviderKeyBundleId, qwEstimatedSize, registrationType);
         ExitOnFailure(hr, "Failed to end session in per-machine process.");
     }
     else
     {
-        hr = RegistrationSessionEnd(&pEngineState->registration, &pEngineState->cache, &pEngineState->variables, &pEngineState->packages, resumeMode, restart, registrationType);
+        hr = RegistrationSessionEnd(&pEngineState->registration, &pEngineState->cache, &pEngineState->variables, &pEngineState->packages, resumeMode, restart, qwEstimatedSize, registrationType);
         ExitOnFailure(hr, "Failed to end session in per-user process.");
     }
 
@@ -880,9 +883,12 @@ extern "C" void ApplyClean(
 static void CalculateKeepRegistration(
     __in BURN_ENGINE_STATE* pEngineState,
     __in BOOL fLog,
-    __inout BOOTSTRAPPER_REGISTRATION_TYPE* pRegistrationType
+    __inout BOOTSTRAPPER_REGISTRATION_TYPE* pRegistrationType,
+    __inout DWORD64* pqwEstimatedSize
     )
 {
+    DWORD64 qwEstimatedSize = pEngineState->section.qwBundleSize;
+
     if (fLog)
     {
         LogId(REPORT_STANDARD, MSG_POST_APPLY_CALCULATE_REGISTRATION);
@@ -911,16 +917,29 @@ static void CalculateKeepRegistration(
         {
             *pRegistrationType = BOOTSTRAPPER_REGISTRATION_TYPE_FULL;
 
-            if (!fLog)
+            if (BURN_PACKAGE_TYPE_MSP == pPackage->type)
             {
-                break;
+                qwEstimatedSize += pPackage->qwSize;
             }
+
+            qwEstimatedSize += pPackage->qwInstallSize;
         }
-        else if (BURN_PACKAGE_REGISTRATION_STATE_PRESENT == pPackage->cacheRegistrationState && BOOTSTRAPPER_REGISTRATION_TYPE_NONE == *pRegistrationType)
+
+        if (BURN_PACKAGE_REGISTRATION_STATE_PRESENT == pPackage->cacheRegistrationState)
         {
-            *pRegistrationType = BOOTSTRAPPER_REGISTRATION_TYPE_INPROGRESS;
+            if (BOOTSTRAPPER_REGISTRATION_TYPE_NONE == *pRegistrationType)
+            {
+                *pRegistrationType = BOOTSTRAPPER_REGISTRATION_TYPE_INPROGRESS;
+            }
+
+            qwEstimatedSize += pPackage->qwSize;
         }
     }
+
+    if (pqwEstimatedSize)
+    {
+        *pqwEstimatedSize = qwEstimatedSize;
+    }
 }
 
 static HRESULT ExecuteDependentRegistrationActions(
diff --git a/src/burn/engine/elevation.cpp b/src/burn/engine/elevation.cpp
index 56a62955..3305cf62 100644
--- a/src/burn/engine/elevation.cpp
+++ b/src/burn/engine/elevation.cpp
@@ -636,6 +636,7 @@ extern "C" HRESULT ElevationSessionEnd(
     __in BURN_RESUME_MODE resumeMode,
     __in BOOTSTRAPPER_APPLY_RESTART restart,
     __in BOOL fDetectedForeignProviderKeyBundleId,
+    __in DWORD64 qwEstimatedSize,
     __in BOOTSTRAPPER_REGISTRATION_TYPE registrationType
     )
 {
@@ -654,6 +655,9 @@ extern "C" HRESULT ElevationSessionEnd(
     hr = BuffWriteNumber(&pbData, &cbData, (DWORD)fDetectedForeignProviderKeyBundleId);
     ExitOnFailure(hr, "Failed to write dependency registration action to message buffer.");
 
+    hr = BuffWriteNumber64(&pbData, &cbData, qwEstimatedSize);
+    ExitOnFailure(hr, "Failed to write estimated size to message buffer.");
+
     hr = BuffWriteNumber(&pbData, &cbData, (DWORD)registrationType);
     ExitOnFailure(hr, "Failed to write registration type to message buffer.");
 
@@ -2580,6 +2584,7 @@ static HRESULT OnSessionEnd(
     SIZE_T iData = 0;
     DWORD dwResumeMode = 0;
     DWORD dwRestart = 0;
+    DWORD64 qwEstimatedSize = 0;
     DWORD dwRegistrationType = 0;
 
     // Deserialize message data.
@@ -2592,11 +2597,14 @@ static HRESULT OnSessionEnd(
     hr = BuffReadNumber(pbData, cbData, &iData, (DWORD*)&pRegistration->fDetectedForeignProviderKeyBundleId);
     ExitOnFailure(hr, "Failed to read dependency registration action.");
 
+    hr = BuffReadNumber64(pbData, cbData, &iData, &qwEstimatedSize);
+    ExitOnFailure(hr, "Failed to read estimated size.");
+
     hr = BuffReadNumber(pbData, cbData, &iData, &dwRegistrationType);
     ExitOnFailure(hr, "Failed to read dependency registration action.");
 
     // suspend session in per-machine process
-    hr = RegistrationSessionEnd(pRegistration, pCache, pVariables, pPackages, (BURN_RESUME_MODE)dwResumeMode, (BOOTSTRAPPER_APPLY_RESTART)dwRestart, (BOOTSTRAPPER_REGISTRATION_TYPE)dwRegistrationType);
+    hr = RegistrationSessionEnd(pRegistration, pCache, pVariables, pPackages, (BURN_RESUME_MODE)dwResumeMode, (BOOTSTRAPPER_APPLY_RESTART)dwRestart, qwEstimatedSize, (BOOTSTRAPPER_REGISTRATION_TYPE)dwRegistrationType);
     ExitOnFailure(hr, "Failed to suspend registration session.");
 
 LExit:
diff --git a/src/burn/engine/elevation.h b/src/burn/engine/elevation.h
index 3484057e..a74b6027 100644
--- a/src/burn/engine/elevation.h
+++ b/src/burn/engine/elevation.h
@@ -44,6 +44,7 @@ HRESULT ElevationSessionEnd(
     __in BURN_RESUME_MODE resumeMode,
     __in BOOTSTRAPPER_APPLY_RESTART restart,
     __in BOOL fDetectedForeignProviderKeyBundleId,
+    __in DWORD64 qwEstimatedSize,
     __in BOOTSTRAPPER_REGISTRATION_TYPE registrationType
     );
 HRESULT ElevationSaveState(
diff --git a/src/burn/engine/plan.cpp b/src/burn/engine/plan.cpp
index 47da22c0..68c1d2ba 100644
--- a/src/burn/engine/plan.cpp
+++ b/src/burn/engine/plan.cpp
@@ -1178,31 +1178,6 @@ extern "C" HRESULT PlanExecutePackage(
         pPackage->rollback = BOOTSTRAPPER_ACTION_STATE_NONE;
     }
 
-    // Add the cache and install size to estimated size if it will be on the machine at the end of the install
-    if (BOOTSTRAPPER_REQUEST_STATE_PRESENT == pPackage->requested ||
-        fRequestedCache ||
-        (BOOTSTRAPPER_PACKAGE_STATE_PRESENT == pPackage->currentState && BOOTSTRAPPER_REQUEST_STATE_ABSENT < pPackage->requested)
-       )
-    {
-        // If the package will remain in the cache, add the package size to the estimated size
-        if (BOOTSTRAPPER_CACHE_TYPE_REMOVE < pPackage->cacheType)
-        {
-            pPlan->qwEstimatedSize += pPackage->qwSize;
-        }
-
-        // If the package will end up installed on the machine, add the install size to the estimated size.
-        if (BOOTSTRAPPER_REQUEST_STATE_CACHE < pPackage->requested)
-        {
-            // MSP packages get cached automatically by windows installer with any embedded cabs, so include that in the size as well
-            if (BURN_PACKAGE_TYPE_MSP == pPackage->type)
-            {
-                pPlan->qwEstimatedSize += pPackage->qwSize;
-            }
-
-            pPlan->qwEstimatedSize += pPackage->qwInstallSize;
-        }
-    }
-
     // Add execute actions.
     switch (pPackage->type)
     {
@@ -3079,7 +3054,6 @@ extern "C" void PlanDump(
     LogStringLine(PlanDumpLevel, "     disallow-removal: %hs", LoggingTrueFalseToString(pPlan->fDisallowRemoval));
     LogStringLine(PlanDumpLevel, "     downgrade: %hs", LoggingTrueFalseToString(pPlan->fDowngrade));
     LogStringLine(PlanDumpLevel, "     registration options: %hs", LoggingRegistrationOptionsToString(pPlan->dwRegistrationOperations));
-    LogStringLine(PlanDumpLevel, "     estimated size: %llu", pPlan->qwEstimatedSize);
     if (pPlan->sczLayoutDirectory)
     {
         LogStringLine(PlanDumpLevel, "     layout directory: %ls", pPlan->sczLayoutDirectory);
diff --git a/src/burn/engine/plan.h b/src/burn/engine/plan.h
index fbb7a5a1..2b04e097 100644
--- a/src/burn/engine/plan.h
+++ b/src/burn/engine/plan.h
@@ -267,8 +267,6 @@ typedef struct _BURN_PLAN
 
     DWORD64 qwCacheSizeTotal;
 
-    DWORD64 qwEstimatedSize;
-
     DWORD cExecutePackagesTotal;
     DWORD cOverallProgressTicksTotal;
 
diff --git a/src/burn/engine/registration.cpp b/src/burn/engine/registration.cpp
index 0ffbc89e..9953be83 100644
--- a/src/burn/engine/registration.cpp
+++ b/src/burn/engine/registration.cpp
@@ -102,6 +102,10 @@ static HRESULT UpdateBundleNameRegistration(
     __in HKEY hkRegistration,
     __in BOOL fInProgressRegistration
     );
+static HRESULT UpdateEstimatedSize(
+    __in HKEY hkRegistration,
+    __in DWORD64 qwEstimatedSize
+    );
 static BOOL IsWuRebootPending();
 static BOOL IsRegistryRebootPending();
 
@@ -600,8 +604,8 @@ extern "C" HRESULT RegistrationSessionBegin(
     )
 {
     HRESULT hr = S_OK;
-    DWORD dwSize = 0;
     HKEY hkRegistration = NULL;
+    BOOL fCreated = FALSE;
     LPWSTR sczPublisher = NULL;
 
     AssertSz(BOOTSTRAPPER_REGISTRATION_TYPE_NONE != registrationType, "Registration type can't be NONE");
@@ -620,7 +624,7 @@ extern "C" HRESULT RegistrationSessionBegin(
     }
 
     // create registration key
-    hr = RegCreate(pRegistration->hkRoot, pRegistration->sczRegistrationKey, KEY_WRITE, &hkRegistration);
+    hr = RegCreateEx(pRegistration->hkRoot, pRegistration->sczRegistrationKey, KEY_WRITE, REG_KEY_DEFAULT, FALSE, NULL, &hkRegistration, &fCreated);
     ExitOnFailure(hr, "Failed to create registration key.");
 
     // Write any ARP values and software tags.
@@ -797,22 +801,12 @@ extern "C" HRESULT RegistrationSessionBegin(
         ExitOnFailure(hr, "Failed to write update registration.");
     }
 
-    // Update estimated size.
-    qwEstimatedSize /= 1024; // Convert bytes to KB
-    if (0 < qwEstimatedSize)
+    // Only set estimated size here for the first time.
+    // It will always get updated at the end of the session.
+    if (fCreated)
     {
-        if (DWORD_MAX < qwEstimatedSize)
-        {
-            // ARP doesn't support QWORDs here
-            dwSize = DWORD_MAX;
-        }
-        else
-        {
-            dwSize = static_cast<DWORD>(qwEstimatedSize);
-        }
-
-        hr = RegWriteNumber(hkRegistration, REGISTRY_BUNDLE_ESTIMATED_SIZE, dwSize);
-        ExitOnFailure(hr, "Failed to write %ls value.", REGISTRY_BUNDLE_ESTIMATED_SIZE);
+        hr = UpdateEstimatedSize(hkRegistration, qwEstimatedSize);
+        ExitOnFailure(hr, "Failed to update estimated size.");
     }
 
     // Register the bundle dependency key.
@@ -879,6 +873,7 @@ extern "C" HRESULT RegistrationSessionEnd(
     __in BURN_PACKAGES* pPackages,
     __in BURN_RESUME_MODE resumeMode,
     __in BOOTSTRAPPER_APPLY_RESTART restart,
+    __in DWORD64 qwEstimatedSize,
     __in BOOTSTRAPPER_REGISTRATION_TYPE registrationType
     )
 {
@@ -921,6 +916,9 @@ extern "C" HRESULT RegistrationSessionEnd(
         // update display name
         hr = UpdateBundleNameRegistration(pRegistration, pVariables, hkRegistration, BOOTSTRAPPER_REGISTRATION_TYPE_INPROGRESS == registrationType);
         ExitOnFailure(hr, "Failed to update name and publisher.");
+
+        hr = UpdateEstimatedSize(hkRegistration, qwEstimatedSize);
+        ExitOnFailure(hr, "Failed to update estimated size.");
     }
 
     // Update resume mode.
@@ -1604,6 +1602,35 @@ LExit:
     return hr;
 }
 
+static HRESULT UpdateEstimatedSize(
+    __in HKEY hkRegistration,
+    __in DWORD64 qwEstimatedSize
+    )
+{
+    HRESULT hr = S_OK;
+    DWORD dwSize = 0;
+
+    qwEstimatedSize /= 1024; // Convert bytes to KB
+    if (0 < qwEstimatedSize)
+    {
+        if (DWORD_MAX < qwEstimatedSize)
+        {
+            // ARP doesn't support QWORDs here
+            dwSize = DWORD_MAX;
+        }
+        else
+        {
+            dwSize = static_cast<DWORD>(qwEstimatedSize);
+        }
+
+        hr = RegWriteNumber(hkRegistration, REGISTRY_BUNDLE_ESTIMATED_SIZE, dwSize);
+        ExitOnFailure(hr, "Failed to write %ls value.", REGISTRY_BUNDLE_ESTIMATED_SIZE);
+    }
+
+LExit:
+    return hr;
+}
+
 static BOOL IsWuRebootPending()
 {
     HRESULT hr = S_OK;
diff --git a/src/burn/engine/registration.h b/src/burn/engine/registration.h
index c07be962..cb5349a3 100644
--- a/src/burn/engine/registration.h
+++ b/src/burn/engine/registration.h
@@ -213,6 +213,7 @@ HRESULT RegistrationSessionEnd(
     __in BURN_PACKAGES* pPackages,
     __in BURN_RESUME_MODE resumeMode,
     __in BOOTSTRAPPER_APPLY_RESTART restart,
+    __in DWORD64 qwEstimatedSize,
     __in BOOTSTRAPPER_REGISTRATION_TYPE registrationType
     );
 HRESULT RegistrationSaveState(
diff --git a/src/burn/test/BurnUnitTest/PlanTest.cpp b/src/burn/test/BurnUnitTest/PlanTest.cpp
index f69606a0..be078c5c 100644
--- a/src/burn/test/BurnUnitTest/PlanTest.cpp
+++ b/src/burn/test/BurnUnitTest/PlanTest.cpp
@@ -106,7 +106,6 @@ namespace Bootstrapper
             ValidateCacheCheckpoint(pPlan, fRollback, dwIndex++, 14);
             Assert::Equal(dwIndex, pPlan->cRollbackCacheActions);
 
-            Assert::Equal(107082ull, pPlan->qwEstimatedSize);
             Assert::Equal(522548ull, pPlan->qwCacheSizeTotal);
 
             fRollback = FALSE;
@@ -252,7 +251,6 @@ namespace Bootstrapper
             dwIndex = 0;
             Assert::Equal(dwIndex, pPlan->cRollbackCacheActions);
 
-            Assert::Equal(0ull, pPlan->qwEstimatedSize);
             Assert::Equal(0ull, pPlan->qwCacheSizeTotal);
 
             fRollback = FALSE;
@@ -392,7 +390,6 @@ namespace Bootstrapper
             dwIndex = 0;
             Assert::Equal(dwIndex, pPlan->cRollbackCacheActions);
 
-            Assert::Equal(18575450ull, pPlan->qwEstimatedSize);
             Assert::Equal(78462280ull, pPlan->qwCacheSizeTotal);
 
             fRollback = FALSE;
@@ -503,7 +500,6 @@ namespace Bootstrapper
             dwIndex = 0;
             Assert::Equal(dwIndex, pPlan->cRollbackCacheActions);
 
-            Assert::Equal(18575450ull, pPlan->qwEstimatedSize);
             Assert::Equal(52254105ull, pPlan->qwCacheSizeTotal);
 
             fRollback = FALSE;
@@ -595,7 +591,6 @@ namespace Bootstrapper
             dwIndex = 0;
             Assert::Equal(dwIndex, pPlan->cRollbackCacheActions);
 
-            Assert::Equal(0ull, pPlan->qwEstimatedSize);
             Assert::Equal(0ull, pPlan->qwCacheSizeTotal);
 
             fRollback = FALSE;
@@ -703,7 +698,6 @@ namespace Bootstrapper
             ValidateCacheCheckpoint(pPlan, fRollback, dwIndex++, 1);
             Assert::Equal(dwIndex, pPlan->cRollbackCacheActions);
 
-            Assert::Equal(35694ull, pPlan->qwEstimatedSize);
             Assert::Equal(168715ull, pPlan->qwCacheSizeTotal);
 
             fRollback = FALSE;
@@ -823,7 +817,6 @@ namespace Bootstrapper
             ValidateCacheCheckpoint(pPlan, fRollback, dwIndex++, 1);
             Assert::Equal(dwIndex, pPlan->cRollbackCacheActions);
 
-            Assert::Equal(35694ull, pPlan->qwEstimatedSize);
             Assert::Equal(168715ull, pPlan->qwCacheSizeTotal);
 
             fRollback = FALSE;
@@ -924,7 +917,6 @@ namespace Bootstrapper
             dwIndex = 0;
             Assert::Equal(dwIndex, pPlan->cRollbackCacheActions);
 
-            Assert::Equal(1463267ull, pPlan->qwEstimatedSize);
             Assert::Equal(119695ull, pPlan->qwCacheSizeTotal);
 
             fRollback = FALSE;
@@ -1021,7 +1013,6 @@ namespace Bootstrapper
             ValidateCacheCheckpoint(pPlan, fRollback, dwIndex++, 1);
             Assert::Equal(dwIndex, pPlan->cRollbackCacheActions);
 
-            Assert::Equal(33743ull, pPlan->qwEstimatedSize);
             Assert::Equal(168715ull, pPlan->qwCacheSizeTotal);
 
             fRollback = FALSE;
@@ -1104,7 +1095,6 @@ namespace Bootstrapper
             dwIndex = 0;
             Assert::Equal(dwIndex, pPlan->cRollbackCacheActions);
 
-            Assert::Equal(0ull, pPlan->qwEstimatedSize);
             Assert::Equal(0ull, pPlan->qwCacheSizeTotal);
 
             fRollback = FALSE;
@@ -1180,7 +1170,6 @@ namespace Bootstrapper
             dwIndex = 0;
             Assert::Equal(dwIndex, pPlan->cRollbackCacheActions);
 
-            Assert::Equal(0ull, pPlan->qwEstimatedSize);
             Assert::Equal(0ull, pPlan->qwCacheSizeTotal);
 
             fRollback = FALSE;
@@ -1273,7 +1262,6 @@ namespace Bootstrapper
             dwIndex = 0;
             Assert::Equal(dwIndex, pPlan->cRollbackCacheActions);
 
-            Assert::Equal(35694ull, pPlan->qwEstimatedSize);
             Assert::Equal(175674ull, pPlan->qwCacheSizeTotal);
 
             fRollback = FALSE;
@@ -1371,7 +1359,6 @@ namespace Bootstrapper
             ValidateCacheCheckpoint(pPlan, fRollback, dwIndex++, 1);
             Assert::Equal(dwIndex, pPlan->cRollbackCacheActions);
 
-            Assert::Equal(35694ull, pPlan->qwEstimatedSize);
             Assert::Equal(168715ull, pPlan->qwCacheSizeTotal);
 
             fRollback = FALSE;
@@ -1473,7 +1460,6 @@ namespace Bootstrapper
             dwIndex = 0;
             Assert::Equal(dwIndex, pPlan->cRollbackCacheActions);
 
-            Assert::Equal(0ull, pPlan->qwEstimatedSize);
             Assert::Equal(0ull, pPlan->qwCacheSizeTotal);
 
             fRollback = FALSE;
@@ -1553,7 +1539,6 @@ namespace Bootstrapper
             dwIndex = 0;
             Assert::Equal(dwIndex, pPlan->cRollbackCacheActions);
 
-            Assert::Equal(0ull, pPlan->qwEstimatedSize);
             Assert::Equal(0ull, pPlan->qwCacheSizeTotal);
 
             fRollback = FALSE;
@@ -1649,7 +1634,6 @@ namespace Bootstrapper
             dwIndex = 0;
             Assert::Equal(dwIndex, pPlan->cRollbackCacheActions);
 
-            Assert::Equal(0ull, pPlan->qwEstimatedSize);
             Assert::Equal(0ull, pPlan->qwCacheSizeTotal);
 
             fRollback = FALSE;
@@ -1720,7 +1704,6 @@ namespace Bootstrapper
             dwIndex = 0;
             Assert::Equal(dwIndex, pPlan->cRollbackCacheActions);
 
-            Assert::Equal(0ull, pPlan->qwEstimatedSize);
             Assert::Equal(0ull, pPlan->qwCacheSizeTotal);
 
             fRollback = FALSE;
@@ -1805,7 +1788,6 @@ namespace Bootstrapper
             dwIndex = 0;
             Assert::Equal(dwIndex, pPlan->cRollbackCacheActions);
 
-            Assert::Equal(0ull, pPlan->qwEstimatedSize);
             Assert::Equal(0ull, pPlan->qwCacheSizeTotal);
 
             fRollback = FALSE;
@@ -1906,7 +1888,6 @@ namespace Bootstrapper
             ValidateCacheCheckpoint(pPlan, fRollback, dwIndex++, 1);
             Assert::Equal(dwIndex, pPlan->cRollbackCacheActions);
 
-            Assert::Equal(56ull, pPlan->qwEstimatedSize);
             Assert::Equal(140ull, pPlan->qwCacheSizeTotal);
 
             fRollback = FALSE;
@@ -2005,7 +1986,6 @@ namespace Bootstrapper
             dwIndex = 0;
             Assert::Equal(dwIndex, pPlan->cRollbackCacheActions);
 
-            Assert::Equal(3055111ull, pPlan->qwEstimatedSize);
             Assert::Equal(6130592ull, pPlan->qwCacheSizeTotal);
 
             fRollback = FALSE;
@@ -2128,7 +2108,6 @@ namespace Bootstrapper
             dwIndex = 0;
             Assert::Equal(dwIndex, pPlan->cRollbackCacheActions);
 
-            Assert::Equal(0ull, pPlan->qwEstimatedSize);
             Assert::Equal(0ull, pPlan->qwCacheSizeTotal);
 
             fRollback = FALSE;
@@ -2245,7 +2224,6 @@ namespace Bootstrapper
             dwIndex = 0;
             Assert::Equal(dwIndex, pPlan->cRollbackCacheActions);
 
-            Assert::Equal(0ull, pPlan->qwEstimatedSize);
             Assert::Equal(0ull, pPlan->qwCacheSizeTotal);
 
             fRollback = FALSE;
@@ -2351,7 +2329,6 @@ namespace Bootstrapper
             ValidateCacheCheckpoint(pPlan, fRollback, dwIndex++, 4);
             Assert::Equal(dwIndex, pPlan->cRollbackCacheActions);
 
-            Assert::Equal(2993671ull, pPlan->qwEstimatedSize);
             Assert::Equal(6048672ull, pPlan->qwCacheSizeTotal);
 
             fRollback = FALSE;
@@ -2451,6 +2428,8 @@ namespace Bootstrapper
                 ReleaseStr(sczFilePath);
             }
 
+            pEngineState->section.qwBundleSize = 1234;
+
             hr = CoreInitializeConstants(pEngineState);
             NativeAssert::Succeeded(hr, "Failed to initialize core constants");
 
diff --git a/src/burn/test/BurnUnitTest/RegistrationTest.cpp b/src/burn/test/BurnUnitTest/RegistrationTest.cpp
index f01d92a4..0075e937 100644
--- a/src/burn/test/BurnUnitTest/RegistrationTest.cpp
+++ b/src/burn/test/BurnUnitTest/RegistrationTest.cpp
@@ -60,6 +60,7 @@ namespace Bootstrapper
             String^ cacheDirectory = Path::Combine(Path::Combine(Environment::GetFolderPath(Environment::SpecialFolder::LocalApplicationData), gcnew String(L"Package Cache")), gcnew String(TEST_BUNDLE_ID));
             String^ cacheExePath = Path::Combine(cacheDirectory, gcnew String(L"setup.exe"));
             DWORD dwRegistrationOptions = BURN_REGISTRATION_ACTION_OPERATIONS_CACHE_BUNDLE;
+            DWORD64 qwEstimatedSize = 1024;
 
             try
             {
@@ -103,7 +104,7 @@ namespace Bootstrapper
                 TestThrowOnFailure(hr, L"Failed to get current process path.");
 
                 // write registration
-                hr = RegistrationSessionBegin(sczCurrentProcess, &registration, &cache, &variables, dwRegistrationOptions, 0, BOOTSTRAPPER_REGISTRATION_TYPE_INPROGRESS);
+                hr = RegistrationSessionBegin(sczCurrentProcess, &registration, &cache, &variables, dwRegistrationOptions, qwEstimatedSize, BOOTSTRAPPER_REGISTRATION_TYPE_INPROGRESS);
                 TestThrowOnFailure(hr, L"Failed to register bundle.");
 
                 // verify that registration was created
@@ -114,7 +115,7 @@ namespace Bootstrapper
                 this->ValidateRunOnceKeyEntry(cacheExePath);
 
                 // end session
-                hr = RegistrationSessionEnd(&registration, &cache, &variables, &packages, BURN_RESUME_MODE_NONE, BOOTSTRAPPER_APPLY_RESTART_NONE, BOOTSTRAPPER_REGISTRATION_TYPE_NONE);
+                hr = RegistrationSessionEnd(&registration, &cache, &variables, &packages, BURN_RESUME_MODE_NONE, BOOTSTRAPPER_APPLY_RESTART_NONE, qwEstimatedSize, BOOTSTRAPPER_REGISTRATION_TYPE_NONE);
                 TestThrowOnFailure(hr, L"Failed to unregister bundle.");
 
                 // verify that registration was removed
@@ -158,6 +159,7 @@ namespace Bootstrapper
             String^ cacheDirectory = Path::Combine(Path::Combine(Environment::GetFolderPath(Environment::SpecialFolder::LocalApplicationData), gcnew String(L"Package Cache")), gcnew String(TEST_BUNDLE_ID));
             String^ cacheExePath = Path::Combine(cacheDirectory, gcnew String(L"setup.exe"));
             DWORD dwRegistrationOptions = 0;
+            DWORD64 qwEstimatedSize = 1024;
             try
             {
                 this->testRegistry->SetUp();
@@ -204,7 +206,7 @@ namespace Bootstrapper
                 //
 
                 // write registration
-                hr = RegistrationSessionBegin(sczCurrentProcess, &registration, &cache, &variables, dwRegistrationOptions, 0, BOOTSTRAPPER_REGISTRATION_TYPE_INPROGRESS);
+                hr = RegistrationSessionBegin(sczCurrentProcess, &registration, &cache, &variables, dwRegistrationOptions, qwEstimatedSize, BOOTSTRAPPER_REGISTRATION_TYPE_INPROGRESS);
                 TestThrowOnFailure(hr, L"Failed to register bundle.");
 
                 // verify that registration was created
@@ -213,7 +215,7 @@ namespace Bootstrapper
                 this->ValidateRunOnceKeyEntry(cacheExePath);
 
                 // complete registration
-                hr = RegistrationSessionEnd(&registration, &cache, &variables, &packages, BURN_RESUME_MODE_ARP, BOOTSTRAPPER_APPLY_RESTART_NONE, BOOTSTRAPPER_REGISTRATION_TYPE_INPROGRESS);
+                hr = RegistrationSessionEnd(&registration, &cache, &variables, &packages, BURN_RESUME_MODE_ARP, BOOTSTRAPPER_APPLY_RESTART_NONE, qwEstimatedSize, BOOTSTRAPPER_REGISTRATION_TYPE_INPROGRESS);
                 TestThrowOnFailure(hr, L"Failed to unregister bundle.");
 
                 // verify that registration was updated
@@ -226,7 +228,7 @@ namespace Bootstrapper
                 //
 
                 // write registration
-                hr = RegistrationSessionBegin(sczCurrentProcess, &registration, &cache, &variables, dwRegistrationOptions, 0, BOOTSTRAPPER_REGISTRATION_TYPE_INPROGRESS);
+                hr = RegistrationSessionBegin(sczCurrentProcess, &registration, &cache, &variables, dwRegistrationOptions, qwEstimatedSize, BOOTSTRAPPER_REGISTRATION_TYPE_INPROGRESS);
                 TestThrowOnFailure(hr, L"Failed to register bundle.");
 
                 // verify that registration was updated
@@ -235,7 +237,7 @@ namespace Bootstrapper
                 this->ValidateRunOnceKeyEntry(cacheExePath);
 
                 // delete registration
-                hr = RegistrationSessionEnd(&registration, &cache, &variables, &packages, BURN_RESUME_MODE_NONE, BOOTSTRAPPER_APPLY_RESTART_NONE, BOOTSTRAPPER_REGISTRATION_TYPE_NONE);
+                hr = RegistrationSessionEnd(&registration, &cache, &variables, &packages, BURN_RESUME_MODE_NONE, BOOTSTRAPPER_APPLY_RESTART_NONE, qwEstimatedSize, BOOTSTRAPPER_REGISTRATION_TYPE_NONE);
                 TestThrowOnFailure(hr, L"Failed to unregister bundle.");
 
                 // verify that registration was removed
@@ -278,6 +280,7 @@ namespace Bootstrapper
             String^ cacheDirectory = Path::Combine(Path::Combine(Environment::GetFolderPath(Environment::SpecialFolder::LocalApplicationData), gcnew String(L"Package Cache")), gcnew String(TEST_BUNDLE_ID));
             String^ cacheExePath = Path::Combine(cacheDirectory, gcnew String(L"setup.exe"));
             DWORD dwRegistrationOptions = 0;
+            DWORD64 qwEstimatedSize = 1024;
             try
             {
                 this->testRegistry->SetUp();
@@ -327,7 +330,7 @@ namespace Bootstrapper
                 //
 
                 // write registration
-                hr = RegistrationSessionBegin(sczCurrentProcess, &registration, &cache, &variables, dwRegistrationOptions, 0, BOOTSTRAPPER_REGISTRATION_TYPE_INPROGRESS);
+                hr = RegistrationSessionBegin(sczCurrentProcess, &registration, &cache, &variables, dwRegistrationOptions, qwEstimatedSize, BOOTSTRAPPER_REGISTRATION_TYPE_INPROGRESS);
                 TestThrowOnFailure(hr, L"Failed to register bundle.");
 
                 // verify that registration was created
@@ -335,7 +338,7 @@ namespace Bootstrapper
                 this->ValidateRunOnceKeyEntry(cacheExePath);
 
                 // complete registration
-                hr = RegistrationSessionEnd(&registration, &cache, &variables, &packages, BURN_RESUME_MODE_ARP, BOOTSTRAPPER_APPLY_RESTART_REQUIRED, BOOTSTRAPPER_REGISTRATION_TYPE_FULL);
+                hr = RegistrationSessionEnd(&registration, &cache, &variables, &packages, BURN_RESUME_MODE_ARP, BOOTSTRAPPER_APPLY_RESTART_REQUIRED, qwEstimatedSize, BOOTSTRAPPER_REGISTRATION_TYPE_FULL);
                 TestThrowOnFailure(hr, L"Failed to unregister bundle.");
 
                 // verify that registration variables were updated
@@ -355,7 +358,7 @@ namespace Bootstrapper
                 //
 
                 // delete registration
-                hr = RegistrationSessionEnd(&registration, &cache, &variables, &packages, BURN_RESUME_MODE_NONE, BOOTSTRAPPER_APPLY_RESTART_NONE, BOOTSTRAPPER_REGISTRATION_TYPE_NONE);
+                hr = RegistrationSessionEnd(&registration, &cache, &variables, &packages, BURN_RESUME_MODE_NONE, BOOTSTRAPPER_APPLY_RESTART_NONE, qwEstimatedSize, BOOTSTRAPPER_REGISTRATION_TYPE_NONE);
                 TestThrowOnFailure(hr, L"Failed to unregister bundle.");
 
                 // verify that registration was removed
@@ -398,6 +401,7 @@ namespace Bootstrapper
             String^ cacheDirectory = Path::Combine(Path::Combine(Environment::GetFolderPath(Environment::SpecialFolder::LocalApplicationData), gcnew String(L"Package Cache")), gcnew String(TEST_BUNDLE_ID));
             String^ cacheExePath = Path::Combine(cacheDirectory, gcnew String(L"setup.exe"));
             DWORD dwRegistrationOptions = 0;
+            DWORD64 qwEstimatedSize = 1024;
             try
             {
                 this->testRegistry->SetUp();
@@ -446,7 +450,7 @@ namespace Bootstrapper
                 //
 
                 // write registration
-                hr = RegistrationSessionBegin(sczCurrentProcess, &registration, &cache, &variables, dwRegistrationOptions, 0, BOOTSTRAPPER_REGISTRATION_TYPE_INPROGRESS);
+                hr = RegistrationSessionBegin(sczCurrentProcess, &registration, &cache, &variables, dwRegistrationOptions, qwEstimatedSize, BOOTSTRAPPER_REGISTRATION_TYPE_INPROGRESS);
                 TestThrowOnFailure(hr, L"Failed to register bundle.");
 
                 // verify that registration was created
@@ -454,7 +458,7 @@ namespace Bootstrapper
                 this->ValidateRunOnceKeyEntry(cacheExePath);
 
                 // finish registration
-                hr = RegistrationSessionEnd(&registration, &cache, &variables, &packages, BURN_RESUME_MODE_ARP, BOOTSTRAPPER_APPLY_RESTART_NONE, BOOTSTRAPPER_REGISTRATION_TYPE_FULL);
+                hr = RegistrationSessionEnd(&registration, &cache, &variables, &packages, BURN_RESUME_MODE_ARP, BOOTSTRAPPER_APPLY_RESTART_NONE, qwEstimatedSize, BOOTSTRAPPER_REGISTRATION_TYPE_FULL);
                 TestThrowOnFailure(hr, L"Failed to register bundle.");
 
                 // verify that registration was updated
@@ -479,7 +483,7 @@ namespace Bootstrapper
                 //
 
                 // write registration
-                hr = RegistrationSessionBegin(sczCurrentProcess, &registration, &cache, &variables, dwRegistrationOptions, 0, BOOTSTRAPPER_REGISTRATION_TYPE_INPROGRESS);
+                hr = RegistrationSessionBegin(sczCurrentProcess, &registration, &cache, &variables, dwRegistrationOptions, qwEstimatedSize, BOOTSTRAPPER_REGISTRATION_TYPE_INPROGRESS);
                 TestThrowOnFailure(hr, L"Failed to register bundle.");
 
                 // verify that registration was updated
@@ -487,7 +491,7 @@ namespace Bootstrapper
                 this->ValidateRunOnceKeyEntry(cacheExePath);
 
                 // delete registration
-                hr = RegistrationSessionEnd(&registration, &cache, &variables, &packages, BURN_RESUME_MODE_NONE, BOOTSTRAPPER_APPLY_RESTART_NONE, BOOTSTRAPPER_REGISTRATION_TYPE_NONE);
+                hr = RegistrationSessionEnd(&registration, &cache, &variables, &packages, BURN_RESUME_MODE_NONE, BOOTSTRAPPER_APPLY_RESTART_NONE, qwEstimatedSize, BOOTSTRAPPER_REGISTRATION_TYPE_NONE);
                 TestThrowOnFailure(hr, L"Failed to unregister bundle.");
 
                 // verify that registration was removed
@@ -533,6 +537,7 @@ namespace Bootstrapper
             BYTE* pbBuffer = NULL;
             SIZE_T cbBuffer = 0;
             DWORD dwRegistrationOptions = 0;
+            DWORD64 qwEstimatedSize = 1024;
             
             String^ cacheDirectory = Path::Combine(Path::Combine(Environment::GetFolderPath(Environment::SpecialFolder::LocalApplicationData), gcnew String(L"Package Cache")), gcnew String(TEST_BUNDLE_ID));
             try
@@ -586,7 +591,7 @@ namespace Bootstrapper
                 TestThrowOnFailure(hr, L"Failed to get current process path.");
 
                 // begin session
-                hr = RegistrationSessionBegin(sczCurrentProcess, &registration, &cache, &variables, dwRegistrationOptions, 0, BOOTSTRAPPER_REGISTRATION_TYPE_INPROGRESS);
+                hr = RegistrationSessionBegin(sczCurrentProcess, &registration, &cache, &variables, dwRegistrationOptions, qwEstimatedSize, BOOTSTRAPPER_REGISTRATION_TYPE_INPROGRESS);
                 TestThrowOnFailure(hr, L"Failed to register bundle.");
 
                 VariableSetNumericHelper(&variables, L"MyBurnVariable1", 42);
@@ -629,7 +634,7 @@ namespace Bootstrapper
                 NativeAssert::StringEqual(L"42", sczValue);
 
                 // end session
-                hr = RegistrationSessionEnd(&registration, &cache, &variables, &packages, BURN_RESUME_MODE_NONE, BOOTSTRAPPER_APPLY_RESTART_NONE, BOOTSTRAPPER_REGISTRATION_TYPE_NONE);
+                hr = RegistrationSessionEnd(&registration, &cache, &variables, &packages, BURN_RESUME_MODE_NONE, BOOTSTRAPPER_APPLY_RESTART_NONE, qwEstimatedSize, BOOTSTRAPPER_REGISTRATION_TYPE_NONE);
                 TestThrowOnFailure(hr, L"Failed to unregister bundle.");
             }
             finally
@@ -673,6 +678,7 @@ namespace Bootstrapper
             String^ cacheDirectory = Path::Combine(Path::Combine(Environment::GetFolderPath(Environment::SpecialFolder::LocalApplicationData), gcnew String(L"Package Cache")), gcnew String(TEST_BUNDLE_ID));
             String^ cacheExePath = Path::Combine(cacheDirectory, gcnew String(L"setup.exe"));
             DWORD dwRegistrationOptions = 0;
+            DWORD64 qwEstimatedSize = 1024;
             try
             {
                 this->testRegistry->SetUp();
@@ -728,7 +734,7 @@ namespace Bootstrapper
                 Assert::Equal((int)BOOTSTRAPPER_RESUME_TYPE_NONE, (int)resumeType);
 
                 // begin session
-                hr = RegistrationSessionBegin(sczCurrentProcess, &registration, &cache, &variables, dwRegistrationOptions, 0, BOOTSTRAPPER_REGISTRATION_TYPE_INPROGRESS);
+                hr = RegistrationSessionBegin(sczCurrentProcess, &registration, &cache, &variables, dwRegistrationOptions, qwEstimatedSize, BOOTSTRAPPER_REGISTRATION_TYPE_INPROGRESS);
                 TestThrowOnFailure(hr, L"Failed to register bundle.");
 
                 VariableSetNumericHelper(&variables, L"MyBurnVariable1", 42);
@@ -767,7 +773,7 @@ namespace Bootstrapper
                 Assert::Equal((int)BOOTSTRAPPER_RESUME_TYPE_INTERRUPTED, (int)resumeType);
 
                 // suspend session
-                hr = RegistrationSessionEnd(&registration, &cache, &variables, &packages, BURN_RESUME_MODE_SUSPEND, BOOTSTRAPPER_APPLY_RESTART_NONE, BOOTSTRAPPER_REGISTRATION_TYPE_INPROGRESS);
+                hr = RegistrationSessionEnd(&registration, &cache, &variables, &packages, BURN_RESUME_MODE_SUSPEND, BOOTSTRAPPER_APPLY_RESTART_NONE, qwEstimatedSize, BOOTSTRAPPER_REGISTRATION_TYPE_INPROGRESS);
                 TestThrowOnFailure(hr, L"Failed to suspend session.");
 
                 // verify that run key was removed
@@ -794,7 +800,7 @@ namespace Bootstrapper
                 this->ValidateRunOnceKeyEntry(cacheExePath);
 
                 // end session
-                hr = RegistrationSessionEnd(&registration, &cache, &variables, &packages, BURN_RESUME_MODE_NONE, BOOTSTRAPPER_APPLY_RESTART_NONE, BOOTSTRAPPER_REGISTRATION_TYPE_NONE);
+                hr = RegistrationSessionEnd(&registration, &cache, &variables, &packages, BURN_RESUME_MODE_NONE, BOOTSTRAPPER_APPLY_RESTART_NONE, qwEstimatedSize, BOOTSTRAPPER_REGISTRATION_TYPE_NONE);
                 TestThrowOnFailure(hr, L"Failed to unregister bundle.");
 
                 // read resume type after session
diff --git a/src/libs/dutil/WixToolset.DUtil/regutil.cpp b/src/libs/dutil/WixToolset.DUtil/regutil.cpp
index 219a6c11..78507ade 100644
--- a/src/libs/dutil/WixToolset.DUtil/regutil.cpp
+++ b/src/libs/dutil/WixToolset.DUtil/regutil.cpp
@@ -107,10 +107,9 @@ DAPI_(HRESULT) RegCreate(
     )
 {
     HRESULT hr = S_OK;
-    DWORD er = ERROR_SUCCESS;
 
-    er = vpfnRegCreateKeyExW(hkRoot, wzSubKey, 0, NULL, REG_OPTION_NON_VOLATILE, dwAccess, NULL, phk, NULL);
-    RegExitOnWin32Error(er, hr, "Failed to create registry key.");
+    hr = RegCreateEx(hkRoot, wzSubKey, dwAccess, REG_KEY_DEFAULT, FALSE, NULL, phk, NULL);
+    RegExitOnFailure(hr, "Failed to create registry key.");
 
 LExit:
     return hr;
diff --git a/src/test/burn/TestBA/TestBA.cs b/src/test/burn/TestBA/TestBA.cs
index fdbcc6d4..b4d74341 100644
--- a/src/test/burn/TestBA/TestBA.cs
+++ b/src/test/burn/TestBA/TestBA.cs
@@ -27,6 +27,7 @@ namespace WixToolset.Test.BA
 
         private string updateBundlePath;
 
+        private bool forceKeepRegistration;
         private bool immediatelyQuit;
         private bool quitAfterDetect;
         private bool explicitlyElevateAndPlanFromOnElevateBegin;
@@ -131,6 +132,12 @@ namespace WixToolset.Test.BA
                 this.explicitlyElevateAndPlanFromOnElevateBegin = false;
             }
 
+            string forceKeepRegistration = this.ReadPackageAction(null, "ForceKeepRegistration");
+            if (String.IsNullOrEmpty(forceKeepRegistration) || !Boolean.TryParse(forceKeepRegistration, out this.forceKeepRegistration))
+            {
+                this.forceKeepRegistration = false;
+            }
+
             string quitAfterDetect = this.ReadPackageAction(null, "QuitAfterDetect");
             if (String.IsNullOrEmpty(quitAfterDetect) || !Boolean.TryParse(quitAfterDetect, out this.quitAfterDetect))
             {
@@ -533,6 +540,16 @@ namespace WixToolset.Test.BA
             this.ShutdownUiThread();
         }
 
+        protected override void OnUnregisterBegin(UnregisterBeginEventArgs args)
+        {
+            if (this.forceKeepRegistration && args.RegistrationType == RegistrationType.None)
+            {
+                args.RegistrationType = RegistrationType.InProgress;
+            }
+
+            this.Log("OnUnregisterBegin, default: {0}, requested: {1}", args.RecommendedRegistrationType, args.RegistrationType);
+        }
+
         private void TestVariables()
         {
             // First make sure we can check and get standard variables of each type.
diff --git a/src/test/burn/WixTestTools/BundleVerifier.cs b/src/test/burn/WixTestTools/BundleVerifier.cs
index 56044c5f..594b19aa 100644
--- a/src/test/burn/WixTestTools/BundleVerifier.cs
+++ b/src/test/burn/WixTestTools/BundleVerifier.cs
@@ -87,7 +87,7 @@ namespace WixTestTools
             }
         }
 
-        public string VerifyRegisteredAndInPackageCache(int? expectedSystemComponent = null)
+        public BundleRegistration VerifyRegisteredAndInPackageCache(int? expectedSystemComponent = null)
         {
             Assert.True(this.TryGetRegistration(out var registration));
 
@@ -99,7 +99,7 @@ namespace WixTestTools
             var expectedCachePath = this.GetExpectedCachedBundlePath();
             WixAssert.StringEqual(expectedCachePath, registration.CachePath, true);
 
-            return registration.CachePath;
+            return registration;
         }
 
         public void VerifyUnregisteredAndRemovedFromPackageCache()
diff --git a/src/test/burn/WixToolsetTest.BurnE2E/BasicFunctionalityTests.cs b/src/test/burn/WixToolsetTest.BurnE2E/BasicFunctionalityTests.cs
index 89c5be9b..efe4f05b 100644
--- a/src/test/burn/WixToolsetTest.BurnE2E/BasicFunctionalityTests.cs
+++ b/src/test/burn/WixToolsetTest.BurnE2E/BasicFunctionalityTests.cs
@@ -82,7 +82,8 @@ namespace WixToolsetTest.BurnE2E
 
             bundle.Install();
 
-            var cachedBundlePath = bundle.VerifyRegisteredAndInPackageCache();
+            var registration = bundle.VerifyRegisteredAndInPackageCache();
+            var cachedBundlePath = registration.CachePath;
 
             // Source file should be installed
             Assert.True(File.Exists(packageSourceCodeInstalled), $"Should have found {packageName} payload installed at: {packageSourceCodeInstalled}");
diff --git a/src/test/burn/WixToolsetTest.BurnE2E/CacheTests.cs b/src/test/burn/WixToolsetTest.BurnE2E/CacheTests.cs
index 943939e3..4d5fb811 100644
--- a/src/test/burn/WixToolsetTest.BurnE2E/CacheTests.cs
+++ b/src/test/burn/WixToolsetTest.BurnE2E/CacheTests.cs
@@ -162,14 +162,14 @@ namespace WixToolsetTest.BurnE2E
                 bundleA.Install(bundleACopiedPath);
             }
 
-            var bundlePackageCachePath = bundleA.VerifyRegisteredAndInPackageCache();
+            var bundlePackageRegistration = bundleA.VerifyRegisteredAndInPackageCache();
 
             packageA.VerifyInstalled(true);
             packageB.VerifyInstalled(false);
 
             testBAController.SetPackageRequestedState("PackageB", RequestState.Present);
 
-            var modifyLogPath = bundleA.Modify(bundlePackageCachePath);
+            var modifyLogPath = bundleA.Modify(bundlePackageRegistration.CachePath);
             bundleA.VerifyRegisteredAndInPackageCache();
 
             packageA.VerifyInstalled(true);
@@ -204,14 +204,14 @@ namespace WixToolsetTest.BurnE2E
 
                 bundleB.Install(bundleBCopiedPath);
 
-                var bundlePackageCachePath = bundleB.VerifyRegisteredAndInPackageCache();
+                var bundlePackageRegistration = bundleB.VerifyRegisteredAndInPackageCache();
 
                 packageA.VerifyInstalled(true);
                 packageB.VerifyInstalled(false);
 
                 testBAController.SetPackageRequestedState("PackageB", RequestState.Present);
 
-                bundleB.Modify(bundlePackageCachePath);
+                bundleB.Modify(bundlePackageRegistration.CachePath);
                 bundleB.VerifyRegisteredAndInPackageCache();
 
                 packageA.VerifyInstalled(true);
diff --git a/src/test/burn/WixToolsetTest.BurnE2E/MsiTransactionTests.cs b/src/test/burn/WixToolsetTest.BurnE2E/MsiTransactionTests.cs
index 3d9748bb..cbfee806 100644
--- a/src/test/burn/WixToolsetTest.BurnE2E/MsiTransactionTests.cs
+++ b/src/test/burn/WixToolsetTest.BurnE2E/MsiTransactionTests.cs
@@ -42,7 +42,7 @@ namespace WixToolsetTest.BurnE2E
 
             bundleAv1.Install();
 
-            var bundleAv1CachedPath = bundleAv1.VerifyRegisteredAndInPackageCache();
+            var bundleAv1Registration = bundleAv1.VerifyRegisteredAndInPackageCache();
 
             // Source file should be installed
             Assert.True(File.Exists(packageASourceCodeInstalled), String.Concat("Should have found Package A payload installed at: ", packageASourceCodeInstalled));
@@ -51,7 +51,7 @@ namespace WixToolsetTest.BurnE2E
 
             bundleAv2.Install();
 
-            var bundleAv2CachedPath = bundleAv2.VerifyRegisteredAndInPackageCache();
+            var bundleAv2Registration = bundleAv2.VerifyRegisteredAndInPackageCache();
 
             // Source file should be upgraded
             Assert.True(File.Exists(packageDSourceCodeInstalled), String.Concat("Should have found Package D payload installed at: ", packageDSourceCodeInstalled));
@@ -61,7 +61,7 @@ namespace WixToolsetTest.BurnE2E
             Assert.False(File.Exists(packageBv1SourceCodeInstalled), String.Concat("Package Bv1 payload should have been removed by upgrade uninstall from: ", packageBv1SourceCodeInstalled));
             Assert.False(File.Exists(packageASourceCodeInstalled), String.Concat("Package A payload should have been removed by upgrade uninstall from: ", packageASourceCodeInstalled));
 
-            bundleAv1.VerifyUnregisteredAndRemovedFromPackageCache(bundleAv1CachedPath);
+            bundleAv1.VerifyUnregisteredAndRemovedFromPackageCache(bundleAv1Registration.CachePath);
 
             // Uninstall everything.
             bundleAv2.Uninstall();
@@ -71,7 +71,7 @@ namespace WixToolsetTest.BurnE2E
             Assert.False(File.Exists(packageBv2SourceCodeInstalled), String.Concat("Package Bv2 payload should have been removed by uninstall from: ", packageBv2SourceCodeInstalled));
             Assert.False(File.Exists(packageCv2SourceCodeInstalled), String.Concat("Package Cv2 payload should have been removed by uninstall from: ", packageCv2SourceCodeInstalled));
 
-            bundleAv2.VerifyUnregisteredAndRemovedFromPackageCache(bundleAv2CachedPath);
+            bundleAv2.VerifyUnregisteredAndRemovedFromPackageCache(bundleAv2Registration.CachePath);
         }
 
         /// <summary>
diff --git a/src/test/burn/WixToolsetTest.BurnE2E/RegistrationTests.cs b/src/test/burn/WixToolsetTest.BurnE2E/RegistrationTests.cs
index 51122c28..01ffa942 100644
--- a/src/test/burn/WixToolsetTest.BurnE2E/RegistrationTests.cs
+++ b/src/test/burn/WixToolsetTest.BurnE2E/RegistrationTests.cs
@@ -3,6 +3,7 @@
 namespace WixToolsetTest.BurnE2E
 {
     using System;
+    using WixToolset.Mba.Core;
     using Xunit;
     using Xunit.Abstractions;
 
@@ -11,8 +12,35 @@ namespace WixToolsetTest.BurnE2E
         public RegistrationTests(ITestOutputHelper testOutputHelper) : base(testOutputHelper) { }
 
         [Fact]
+        public void AllowsBAToKeepRegistration()
+        {
+            this.CreatePackageInstaller("PackageA");
+            var bundleA = this.CreateBundleInstaller("BundleA");
+            var testBAController = this.CreateTestBAController();
+
+            testBAController.SetPackageRequestedState("PackageA", RequestState.Absent);
+            testBAController.SetForceKeepRegistration();
+
+            bundleA.Install();
+            var initialRegistration = bundleA.VerifyRegisteredAndInPackageCache();
+
+            Assert.NotNull(initialRegistration.EstimatedSize);
+
+            testBAController.SetForceKeepRegistration(null);
+            testBAController.ResetPackageStates("PackageA");
+
+            bundleA.Install();
+            var finalRegistration = bundleA.VerifyRegisteredAndInPackageCache();
+
+            // Verifies https://github.com/wixtoolset/issues/issues/4039
+            Assert.NotNull(finalRegistration.EstimatedSize);
+            Assert.InRange(finalRegistration.EstimatedSize.Value, initialRegistration.EstimatedSize.Value + 1, Int32.MaxValue);
+        }
+
+        [Fact]
         public void AutomaticallyUncachesBundleWhenNotInstalled()
         {
+            this.CreatePackageInstaller("PackageA");
             var bundleA = this.CreateBundleInstaller("BundleA");
             var testBAController = this.CreateTestBAController();
 
@@ -40,6 +68,7 @@ namespace WixToolsetTest.BurnE2E
         [Fact]
         public void RegistersInARPIfPrecached()
         {
+            this.CreatePackageInstaller("PackageA");
             var bundleA = this.CreateBundleInstaller("BundleA");
 
             bundleA.ManuallyCache();
diff --git a/src/test/burn/WixToolsetTest.BurnE2E/TestBAController.cs b/src/test/burn/WixToolsetTest.BurnE2E/TestBAController.cs
index fa553919..8e6611a2 100644
--- a/src/test/burn/WixToolsetTest.BurnE2E/TestBAController.cs
+++ b/src/test/burn/WixToolsetTest.BurnE2E/TestBAController.cs
@@ -47,6 +47,11 @@ namespace WixToolsetTest.BurnE2E
             this.SetBurnTestValue("ExplicitlyElevateAndPlanFromOnElevateBegin", value);
         }
 
+        public void SetForceKeepRegistration(string value = "true")
+        {
+            this.SetBurnTestValue("ForceKeepRegistration", value);
+        }
+
         public void SetImmediatelyQuit(string value = "true")
         {
             this.SetBurnTestValue("ImmediatelyQuit", value);