4 files changed, 33 insertions, 10 deletions

diff --git a/src/engine/apply.cpp b/src/engine/apply.cpp
index 5402e55c..f80baf26 100644
--- a/src/engine/apply.cpp
+++ b/src/engine/apply.cpp
@@ -869,6 +869,12 @@ static HRESULT ApplyExtractContainer(
 {
     HRESULT hr = S_OK;
 
+    if (pContainer->qwCommittedCacheProgress)
+    {
+        pContext->qwSuccessfulCacheProgress -= pContainer->qwCommittedCacheProgress;
+        pContainer->qwCommittedCacheProgress = 0;
+    }
+
     if (!pContainer->fActuallyAttached)
     {
         hr = ApplyAcquireContainerOrPayload(pContext, pContainer, NULL, NULL);
@@ -884,7 +890,8 @@ static HRESULT ApplyExtractContainer(
         CacheSetLastUsedSource(pContext->pVariables, pContext->sczLastUsedFolderCandidate, pContainer->sczFilePath);
     }
 
-    pContext->qwSuccessfulCacheProgress += pContainer->qwFileSize;
+    pContext->qwSuccessfulCacheProgress += pContainer->qwExtractSizeTotal;
+    pContainer->qwCommittedCacheProgress += pContainer->qwExtractSizeTotal;
 
 LExit:
     ReleaseNullStr(pContext->sczLastUsedFolderCandidate);
@@ -1110,7 +1117,7 @@ static HRESULT ExtractContainer(
         for (DWORD iExtract = 0; iExtract < pContext->pPayloads->cPayloads; ++iExtract)
         {
             BURN_PAYLOAD* pExtract = pContext->pPayloads->rgPayloads + iExtract;
-            if (pExtract->sczUnverifiedPath && CSTR_EQUAL == ::CompareStringW(LOCALE_INVARIANT, 0, sczExtractPayloadId, -1, pExtract->sczSourcePath, -1))
+            if (pExtract->sczUnverifiedPath && pExtract->cRemainingInstances && CSTR_EQUAL == ::CompareStringW(LOCALE_INVARIANT, 0, sczExtractPayloadId, -1, pExtract->sczSourcePath, -1))
             {
                 hr = PreparePayloadDestinationPath(pExtract->sczUnverifiedPath);
                 ExitOnFailure(hr, "Failed to prepare payload destination path: %ls", pExtract->sczUnverifiedPath);
@@ -1747,7 +1754,11 @@ static HRESULT CompleteCacheProgress(
     if (PROGRESS_CONTINUE == dwResult)
     {
         pContext->pCacheContext->qwSuccessfulCacheProgress += qwFileSize;
-        if (pContext->pPayloadGroupItem)
+        if (pContext->pContainer)
+        {
+            pContext->pContainer->qwCommittedCacheProgress += qwFileSize;
+        }
+        else if (pContext->pPayloadGroupItem)
         {
             pContext->pPayloadGroupItem->qwCommittedCacheProgress += qwFileSize;
         }
diff --git a/src/engine/container.h b/src/engine/container.h
index d38016cc..7c5c2b5f 100644
--- a/src/engine/container.h
+++ b/src/engine/container.h
@@ -76,6 +76,8 @@ typedef struct _BURN_CONTAINER
     BOOL fPlanned;
     LPWSTR sczSourcePath;
     LPWSTR sczUnverifiedPath;
+    DWORD64 qwExtractSizeTotal;
+    DWORD64 qwCommittedCacheProgress;
 } BURN_CONTAINER;
 
 typedef struct _BURN_CONTAINERS
diff --git a/src/engine/plan.cpp b/src/engine/plan.cpp
index 55fe7ddf..8421d87b 100644
--- a/src/engine/plan.cpp
+++ b/src/engine/plan.cpp
@@ -431,6 +431,7 @@ extern "C" HRESULT PlanLayoutBundle(
     pCacheAction->bundleLayout.qwBundleSize = qwBundleSize;
     pCacheAction->bundleLayout.pPayloadGroup = pLayoutPayloads;
 
+    // Acquire + Verify
     pPlan->qwCacheSizeTotal += 2 * qwBundleSize;
 
     ++pPlan->cOverallProgressTicksTotal;
@@ -1005,6 +1006,7 @@ extern "C" HRESULT PlanLayoutContainer(
             pCacheAction->type = BURN_CACHE_ACTION_TYPE_CONTAINER;
             pCacheAction->container.pContainer = pContainer;
 
+            // Acquire + Verify
             pPlan->qwCacheSizeTotal += 2 * pContainer->qwFileSize;
         }
     }
@@ -1012,11 +1014,9 @@ extern "C" HRESULT PlanLayoutContainer(
     {
         if (!pContainer->fActuallyAttached)
         {
+            // Acquire
             pPlan->qwCacheSizeTotal += pContainer->qwFileSize;
         }
-
-        // TODO: This should be the sum of all uncompressed payloads in the container, ideally restricted to the payloads that were actually planned.
-        pPlan->qwCacheSizeTotal += pContainer->qwFileSize;
     }
 
     if (!pContainer->sczUnverifiedPath)
@@ -1829,6 +1829,8 @@ static void ResetPlannedContainerState(
     )
 {
     pContainer->fPlanned = FALSE;
+    pContainer->qwExtractSizeTotal = 0;
+    pContainer->qwCommittedCacheProgress = 0;
 }
 
 static void ResetPlannedPayloadsState(
@@ -2246,9 +2248,17 @@ static HRESULT ProcessPayloadGroup(
 
         if (!pPlan->sczLayoutDirectory || !pPayload->pContainer)
         {
+            // Acquire + Verify
             pPlan->qwCacheSizeTotal += 2 * pPayload->qwFileSize;
         }
 
+        if (!pPlan->sczLayoutDirectory && pPayload->pContainer && 1 == pPayload->cRemainingInstances)
+        {
+            // Extract
+            pPlan->qwCacheSizeTotal += pPayload->qwFileSize;
+            pPayload->pContainer->qwExtractSizeTotal += pPayload->qwFileSize;
+        }
+
         if (!pPayload->sczUnverifiedPath)
         {
             hr = CacheCalculatePayloadWorkingPath(pPlan->wzBundleId, pPayload, &pPayload->sczUnverifiedPath);
diff --git a/src/test/BurnUnitTest/PlanTest.cpp b/src/test/BurnUnitTest/PlanTest.cpp
index aa9deaf6..c073696b 100644
--- a/src/test/BurnUnitTest/PlanTest.cpp
+++ b/src/test/BurnUnitTest/PlanTest.cpp
@@ -71,7 +71,7 @@ namespace Bootstrapper
             Assert::Equal(dwIndex, pPlan->cRollbackCacheActions);
 
             Assert::Equal(107082ull, pPlan->qwEstimatedSize);
-            Assert::Equal(202458ull, pPlan->qwCacheSizeTotal);
+            Assert::Equal(303687ull, pPlan->qwCacheSizeTotal);
 
             fRollback = FALSE;
             dwIndex = 0;
@@ -308,7 +308,7 @@ namespace Bootstrapper
             Assert::Equal(dwIndex, pPlan->cRollbackCacheActions);
 
             Assert::Equal(35694ull, pPlan->qwEstimatedSize);
-            Assert::Equal(67486ull, pPlan->qwCacheSizeTotal);
+            Assert::Equal(101229ull, pPlan->qwCacheSizeTotal);
 
             fRollback = FALSE;
             dwIndex = 0;
@@ -388,7 +388,7 @@ namespace Bootstrapper
             Assert::Equal(dwIndex, pPlan->cRollbackCacheActions);
 
             Assert::Equal(33743ull, pPlan->qwEstimatedSize);
-            Assert::Equal(67486ull, pPlan->qwCacheSizeTotal);
+            Assert::Equal(101229ull, pPlan->qwCacheSizeTotal);
 
             fRollback = FALSE;
             dwIndex = 0;
@@ -458,7 +458,7 @@ namespace Bootstrapper
             Assert::Equal(dwIndex, pPlan->cRollbackCacheActions);
 
             Assert::Equal(35694ull, pPlan->qwEstimatedSize);
-            Assert::Equal(67486ull, pPlan->qwCacheSizeTotal);
+            Assert::Equal(101229ull, pPlan->qwCacheSizeTotal);
 
             fRollback = FALSE;
             dwIndex = 0;