Make the estimated size in ARP a little more accurate.

Fixes 4039

9 files changed, 112 insertions, 99 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