Fix issue where applying rollback file would incorrectly uninstall workload packs (#38737)

dsplaisted · web-flow · commit 17bc1dda7ee8 · 2024-02-13T21:29:11.000-05:00
diff --git a/src/Cli/dotnet/commands/dotnet-workload/install/WorkloadInstallCommand.cs b/src/Cli/dotnet/commands/dotnet-workload/install/WorkloadInstallCommand.cs
@@ -216,6 +216,11 @@ private void InstallWorkloadsWithInstallRecord(
                         installer.InstallWorkloadManifest(manifestUpdate, context, offlineCache, rollback);
                     }
 
+                    if (usingRollback)
+                    {
+                        UpdateInstallState(true, manifestsToUpdate);
+                    }
+
                     _workloadResolver.RefreshWorkloadManifests();
 
                     installer.InstallWorkloads(workloadIds, sdkFeatureBand, context, offlineCache);
@@ -227,10 +232,6 @@ private void InstallWorkloadsWithInstallRecord(
                         recordRepo.WriteWorkloadInstallationRecord(workloadId, sdkFeatureBand);
                     }
 
-                    if (usingRollback)
-                    {
-                        UpdateInstallState(true, manifestsToUpdate);
-                    }
                 },
                 rollback: () =>
                 {
@@ -242,6 +243,9 @@ private void InstallWorkloadsWithInstallRecord(
                         installer.GetWorkloadInstallationRecordRepository()
                             .DeleteWorkloadInstallationRecord(workloadId, sdkFeatureBand);
                     }
+
+                    //  Refresh the workload manifests to make sure that the resolver has the updated state after the rollback
+                    _workloadResolver.RefreshWorkloadManifests();
                 });
 
         }
diff --git a/src/Cli/dotnet/commands/dotnet-workload/update/WorkloadUpdateCommand.cs b/src/Cli/dotnet/commands/dotnet-workload/update/WorkloadUpdateCommand.cs
@@ -150,17 +150,19 @@ private void UpdateWorkloadsWithInstallRecord(
                         _workloadInstaller.InstallWorkloadManifest(manifestUpdate, context, offlineCache, rollback);
                     }
 
+					UpdateInstallState(useRollback, manifestsToUpdate);
+
                     _workloadResolver.RefreshWorkloadManifests();
 
                     var workloads = GetUpdatableWorkloads();
 
                     _workloadInstaller.InstallWorkloads(workloads, sdkFeatureBand, context, offlineCache);
-
-                    UpdateInstallState(useRollback, manifestsToUpdate);
                 },
                 rollback: () =>
                 {
                     //  Nothing to roll back at this level, InstallWorkloadManifest and InstallWorkloadPacks handle the transaction rollback
+                    //  We will refresh the workload manifests to make sure that the resolver has the updated state after the rollback
+                    _workloadResolver.RefreshWorkloadManifests();
                 });
         }
 
diff --git a/src/Resolvers/Microsoft.NET.Sdk.WorkloadManifestReader/IWorkloadManifestProvider.cs b/src/Resolvers/Microsoft.NET.Sdk.WorkloadManifestReader/IWorkloadManifestProvider.cs
@@ -9,6 +9,7 @@ namespace Microsoft.NET.Sdk.WorkloadManifestReader
     /// </summary>
     public interface IWorkloadManifestProvider
     {
+        void RefreshWorkloadManifests();
         IEnumerable<ReadableWorkloadManifest> GetManifests();
 
         string GetSdkFeatureBand();
diff --git a/src/Resolvers/Microsoft.NET.Sdk.WorkloadManifestReader/SdkDirectoryWorkloadManifestProvider.cs b/src/Resolvers/Microsoft.NET.Sdk.WorkloadManifestReader/SdkDirectoryWorkloadManifestProvider.cs
@@ -25,9 +25,12 @@ public partial class SdkDirectoryWorkloadManifestProvider : IWorkloadManifestPro
             "microsoft.net.workload.maccatalyst", "microsoft.net.workload.macos", "microsoft.net.workload.tvos", "microsoft.net.workload.mono.toolchain" };
         private readonly Dictionary<string, int>? _knownManifestIdsAndOrder;
 
-        private readonly WorkloadSet? _workloadSet;
-        private readonly WorkloadSet? _manifestsFromInstallState;
-        private readonly string? _installStateFilePath;
+        private readonly string? _workloadSetVersionFromConstructor;
+        private readonly string? _globalJsonPathFromConstructor;
+
+        private WorkloadSet? _workloadSet;
+        private WorkloadSet? _manifestsFromInstallState;
+        private string? _installStateFilePath;
 
         public SdkDirectoryWorkloadManifestProvider(string sdkRootPath, string sdkVersion, string? userProfileDir, string? globalJsonPath)
             : this(sdkRootPath, sdkVersion, Environment.GetEnvironmentVariable, userProfileDir, globalJsonPath)
@@ -59,6 +62,8 @@ internal SdkDirectoryWorkloadManifestProvider(string sdkRootPath, string sdkVers
 
             _sdkRootPath = sdkRootPath;
             _sdkVersionBand = new SdkFeatureBand(sdkVersion);
+            _workloadSetVersionFromConstructor = workloadSetVersion;
+            _globalJsonPathFromConstructor = globalJsonPath;
 
             var knownManifestIdsFilePath = Path.Combine(_sdkRootPath, "sdk", sdkVersion, "KnownWorkloadManifests.txt");
             if (!File.Exists(knownManifestIdsFilePath))
@@ -102,23 +107,28 @@ internal SdkDirectoryWorkloadManifestProvider(string sdkRootPath, string sdkVers
 
             _manifestRoots ??= Array.Empty<string>();
 
+            RefreshWorkloadManifests();
+        }
+
+        public void RefreshWorkloadManifests()
+        {
             var availableWorkloadSets = GetAvailableWorkloadSets();
 
-            if (workloadSetVersion != null)
+            if (_workloadSetVersionFromConstructor != null)
             {
-                if (!availableWorkloadSets.TryGetValue(workloadSetVersion, out _workloadSet))
+                if (!availableWorkloadSets.TryGetValue(_workloadSetVersionFromConstructor, out _workloadSet))
                 {
-                    throw new FileNotFoundException(string.Format(Strings.WorkloadVersionNotFound, workloadSetVersion));
+                    throw new FileNotFoundException(string.Format(Strings.WorkloadVersionNotFound, _workloadSetVersionFromConstructor));
                 }
             }
             else
             {
-                string? globalJsonWorkloadSetVersion = GlobalJsonReader.GetWorkloadVersionFromGlobalJson(globalJsonPath);
+                string? globalJsonWorkloadSetVersion = GlobalJsonReader.GetWorkloadVersionFromGlobalJson(_globalJsonPathFromConstructor);
                 if (globalJsonWorkloadSetVersion != null)
                 {
                     if (!availableWorkloadSets.TryGetValue(globalJsonWorkloadSetVersion, out _workloadSet))
                     {
-                        throw new FileNotFoundException(string.Format(Strings.WorkloadVersionFromGlobalJsonNotFound, globalJsonWorkloadSetVersion, globalJsonPath));
+                        throw new FileNotFoundException(string.Format(Strings.WorkloadVersionFromGlobalJsonNotFound, globalJsonWorkloadSetVersion, _globalJsonPathFromConstructor));
                     }
                 }
                 else
diff --git a/src/Resolvers/Microsoft.NET.Sdk.WorkloadManifestReader/TempDirectoryWorkloadManifestProvider.cs b/src/Resolvers/Microsoft.NET.Sdk.WorkloadManifestReader/TempDirectoryWorkloadManifestProvider.cs
@@ -14,6 +14,8 @@ public TempDirectoryWorkloadManifestProvider(string manifestsPath, string sdkFea
             _sdkVersionBand = sdkFeatureBand;
         }
 
+        public void RefreshWorkloadManifests() { }
+
         public IEnumerable<ReadableWorkloadManifest>
             GetManifests()
         {
diff --git a/src/Resolvers/Microsoft.NET.Sdk.WorkloadManifestReader/WorkloadResolver.cs b/src/Resolvers/Microsoft.NET.Sdk.WorkloadManifestReader/WorkloadResolver.cs
@@ -102,6 +102,8 @@ public void RefreshWorkloadManifests()
             {
                 throw new InvalidOperationException("Resolver was created without provider and cannot be refreshed");
             }
+
+            _manifestProvider.RefreshWorkloadManifests();
             _manifests.Clear();
             LoadManifestsFromProvider(_manifestProvider);
             ComposeWorkloadManifests();
@@ -735,6 +737,7 @@ public EmptyWorkloadManifestProvider(string sdkFeatureBand)
                 _sdkFeatureBand = sdkFeatureBand;
             }
 
+            public void RefreshWorkloadManifests() { }
             public Dictionary<string, WorkloadSet> GetAvailableWorkloadSets() => new();
             public IEnumerable<ReadableWorkloadManifest> GetManifests() => Enumerable.Empty<ReadableWorkloadManifest>();
             public string GetSdkFeatureBand() => _sdkFeatureBand;
diff --git a/src/Tests/Microsoft.NET.Sdk.WorkloadManifestReader.Tests/FakeManifestProvider.cs b/src/Tests/Microsoft.NET.Sdk.WorkloadManifestReader.Tests/FakeManifestProvider.cs
@@ -20,6 +20,8 @@ public FakeManifestProvider(params (string manifest, string? localizationCatalog
             _filePaths = filePaths;
         }
 
+        public void RefreshWorkloadManifests() { }
+
         public IEnumerable<ReadableWorkloadManifest> GetManifests()
         {
             foreach (var filePath in _filePaths)
@@ -46,6 +48,8 @@ internal class InMemoryFakeManifestProvider : IWorkloadManifestProvider, IEnumer
 
         public void Add(string id, string content) => _manifests.Add((id, Encoding.UTF8.GetBytes(content)));
 
+        public void RefreshWorkloadManifests() { }
+
         public IEnumerable<ReadableWorkloadManifest> GetManifests()
             => _manifests.Select(m => new ReadableWorkloadManifest(
                 m.id,
diff --git a/src/Tests/dotnet-workload-install.Tests/MockManifestProvider.cs b/src/Tests/dotnet-workload-install.Tests/MockManifestProvider.cs
@@ -29,6 +29,8 @@ public MockManifestProvider(params (string name, string path, string featureBand
 
         public Dictionary<string, WorkloadSet> GetAvailableWorkloadSets() => new();
 
+        public void RefreshWorkloadManifests() { }
+
         public IEnumerable<ReadableWorkloadManifest> GetManifests()
             {
                 foreach ((var id, var path, var featureBand) in _manifests)

Original file line number	Diff line number	Diff line change
`@@ -216,6 +216,11 @@ private void InstallWorkloadsWithInstallRecord(`
`216`	`216`	`installer.InstallWorkloadManifest(manifestUpdate, context, offlineCache, rollback);`
`217`	`217`	`}`
`218`	`218`
	`219`	`+ if (usingRollback)`
	`220`	`+ {`
	`221`	`+ UpdateInstallState(true, manifestsToUpdate);`
	`222`	`+ }`
	`223`	`+`
`219`	`224`	`_workloadResolver.RefreshWorkloadManifests();`
`220`	`225`
`221`	`226`	`installer.InstallWorkloads(workloadIds, sdkFeatureBand, context, offlineCache);`
`@@ -227,10 +232,6 @@ private void InstallWorkloadsWithInstallRecord(`
`227`	`232`	`recordRepo.WriteWorkloadInstallationRecord(workloadId, sdkFeatureBand);`
`228`	`233`	`}`
`229`	`234`
`230`		`- if (usingRollback)`
`231`		`- {`
`232`		`- UpdateInstallState(true, manifestsToUpdate);`
`233`		`- }`
`234`	`235`	`},`
`235`	`236`	`rollback: () =>`
`236`	`237`	`{`
`@@ -242,6 +243,9 @@ private void InstallWorkloadsWithInstallRecord(`
`242`	`243`	`installer.GetWorkloadInstallationRecordRepository()`
`243`	`244`	`.DeleteWorkloadInstallationRecord(workloadId, sdkFeatureBand);`
`244`	`245`	`}`
	`246`	`+`
	`247`	`+ // Refresh the workload manifests to make sure that the resolver has the updated state after the rollback`
	`248`	`+ _workloadResolver.RefreshWorkloadManifests();`
`245`	`249`	`});`
`246`	`250`
`247`	`251`	`}`
Original file line number	Diff line number	Diff line change
`@@ -150,17 +150,19 @@ private void UpdateWorkloadsWithInstallRecord(`
`150`	`150`	`_workloadInstaller.InstallWorkloadManifest(manifestUpdate, context, offlineCache, rollback);`
`151`	`151`	`}`
`152`	`152`
	`153`	`+ UpdateInstallState(useRollback, manifestsToUpdate);`
	`154`	`+`
`153`	`155`	`_workloadResolver.RefreshWorkloadManifests();`
`154`	`156`
`155`	`157`	`var workloads = GetUpdatableWorkloads();`
`156`	`158`
`157`	`159`	`_workloadInstaller.InstallWorkloads(workloads, sdkFeatureBand, context, offlineCache);`
`158`		`-`
`159`		`- UpdateInstallState(useRollback, manifestsToUpdate);`
`160`	`160`	`},`
`161`	`161`	`rollback: () =>`
`162`	`162`	`{`
`163`	`163`	`// Nothing to roll back at this level, InstallWorkloadManifest and InstallWorkloadPacks handle the transaction rollback`
	`164`	`+ // We will refresh the workload manifests to make sure that the resolver has the updated state after the rollback`
	`165`	`+ _workloadResolver.RefreshWorkloadManifests();`
`164`	`166`	`});`
`165`	`167`	`}`
`166`	`168`
Original file line number	Diff line number	Diff line change
`@@ -9,6 +9,7 @@ namespace Microsoft.NET.Sdk.WorkloadManifestReader`
`9`	`9`	`/// </summary>`
`10`	`10`	`public interface IWorkloadManifestProvider`
`11`	`11`	`{`
	`12`	`+ void RefreshWorkloadManifests();`
`12`	`13`	`IEnumerable<ReadableWorkloadManifest> GetManifests();`
`13`	`14`
`14`	`15`	`string GetSdkFeatureBand();`
Original file line number	Diff line number	Diff line change
`@@ -14,6 +14,8 @@ public TempDirectoryWorkloadManifestProvider(string manifestsPath, string sdkFea`
`14`	`14`	`_sdkVersionBand = sdkFeatureBand;`
`15`	`15`	`}`
`16`	`16`
	`17`	`+ public void RefreshWorkloadManifests() { }`
	`18`	`+`
`17`	`19`	`public IEnumerable<ReadableWorkloadManifest>`
`18`	`20`	`GetManifests()`
`19`	`21`	`{`
Original file line number	Diff line number	Diff line change
`@@ -102,6 +102,8 @@ public void RefreshWorkloadManifests()`
`102`	`102`	`{`
`103`	`103`	`throw new InvalidOperationException("Resolver was created without provider and cannot be refreshed");`
`104`	`104`	`}`
	`105`	`+`
	`106`	`+ _manifestProvider.RefreshWorkloadManifests();`
`105`	`107`	`_manifests.Clear();`
`106`	`108`	`LoadManifestsFromProvider(_manifestProvider);`
`107`	`109`	`ComposeWorkloadManifests();`
`@@ -735,6 +737,7 @@ public EmptyWorkloadManifestProvider(string sdkFeatureBand)`
`735`	`737`	`_sdkFeatureBand = sdkFeatureBand;`
`736`	`738`	`}`
`737`	`739`
	`740`	`+ public void RefreshWorkloadManifests() { }`
`738`	`741`	`public Dictionary<string, WorkloadSet> GetAvailableWorkloadSets() => new();`
`739`	`742`	`public IEnumerable<ReadableWorkloadManifest> GetManifests() => Enumerable.Empty<ReadableWorkloadManifest>();`
`740`	`743`	`public string GetSdkFeatureBand() => _sdkFeatureBand;`
Original file line number	Diff line number	Diff line change
`@@ -20,6 +20,8 @@ public FakeManifestProvider(params (string manifest, string? localizationCatalog`
`20`	`20`	`_filePaths = filePaths;`
`21`	`21`	`}`
`22`	`22`
	`23`	`+ public void RefreshWorkloadManifests() { }`
	`24`	`+`
`23`	`25`	`public IEnumerable<ReadableWorkloadManifest> GetManifests()`
`24`	`26`	`{`
`25`	`27`	`foreach (var filePath in _filePaths)`
`@@ -46,6 +48,8 @@ internal class InMemoryFakeManifestProvider : IWorkloadManifestProvider, IEnumer`
`46`	`48`
`47`	`49`	`public void Add(string id, string content) => _manifests.Add((id, Encoding.UTF8.GetBytes(content)));`
`48`	`50`
	`51`	`+ public void RefreshWorkloadManifests() { }`
	`52`	`+`
`49`	`53`	`public IEnumerable<ReadableWorkloadManifest> GetManifests()`
`50`	`54`	`=> _manifests.Select(m => new ReadableWorkloadManifest(`
`51`	`55`	`m.id,`
Original file line number	Diff line number	Diff line change
`@@ -29,6 +29,8 @@ public MockManifestProvider(params (string name, string path, string featureBand`
`29`	`29`
`30`	`30`	`public Dictionary<string, WorkloadSet> GetAvailableWorkloadSets() => new();`
`31`	`31`
	`32`	`+ public void RefreshWorkloadManifests() { }`
	`33`	`+`
`32`	`34`	`public IEnumerable<ReadableWorkloadManifest> GetManifests()`
`33`	`35`	`{`
`34`	`36`	`foreach ((var id, var path, var featureBand) in _manifests)`