Refactor workload set and workload install/update logic

src/Cli/Microsoft.DotNet.Cli.Utils/PathUtility.cs

@@ -102,7 +102,7 @@ public static bool TryDeleteDirectory(string directoryPath)
         /// and continues to its parent until it fails. Returns whether it succeeded
         /// in deleting the file it was intended to delete.
         /// </summary>
-        public static bool DeleteFileAndEmptyParents(string path)
+        public static bool DeleteFileAndEmptyParents(string path, int maxDirectoriesToDelete = int.MaxValue)
         {
             if (!File.Exists(path))
             {
@@ -112,9 +112,13 @@ public static bool DeleteFileAndEmptyParents(string path)
             File.Delete(path);
             var dir = Path.GetDirectoryName(path);
 
-            while (!Directory.EnumerateFileSystemEntries(dir).Any())
+            int directoriesDeleted = 0;
+
+            while (!Directory.EnumerateFileSystemEntries(dir).Any() &&
+                directoriesDeleted < maxDirectoriesToDelete)
             {
                 Directory.Delete(dir);
+                directoriesDeleted++;
                 dir = Path.GetDirectoryName(dir);
             }
 

src/Cli/dotnet/Installer/Windows/InstallMessageDispatcher.cs

@@ -183,7 +183,7 @@ public InstallResponseMessage SendSaveInstallStateManifestVersions(SdkFeatureBan
         /// <param name="sdkFeatureBand">The SDK feature band of the install state file to write</param>
         /// <param name="newMode">Whether to use workload sets or not</param>
         /// <returns></returns>
-        public InstallResponseMessage SendUpdateWorkloadModeRequest(SdkFeatureBand sdkFeatureBand, bool newMode)
+        public InstallResponseMessage SendUpdateWorkloadModeRequest(SdkFeatureBand sdkFeatureBand, bool? newMode)
         {
             return Send(new InstallRequestMessage
             {

src/Cli/dotnet/Installer/Windows/InstallRequestMessage.cs

@@ -123,7 +123,7 @@ public string WorkloadId
         /// <summary>
         /// The new mode to use: workloadset or loosemanifests
         /// </summary>
-        public bool UseWorkloadSets
+        public bool? UseWorkloadSets
         {
             get; set;
         }

src/Cli/dotnet/commands/InstallingWorkloadCommand.cs

@@ -14,6 +14,7 @@
 using Microsoft.DotNet.Cli.Utils;
 using Microsoft.DotNet.ToolPackage;
 using Microsoft.DotNet.Workloads.Workload.Install;
+using Microsoft.DotNet.Workloads.Workload.Update;
 using Microsoft.Extensions.EnvironmentAbstractions;
 using Microsoft.NET.Sdk.WorkloadManifestReader;
 using NuGet.Versioning;
@@ -34,7 +35,8 @@ internal abstract class InstallingWorkloadCommand : WorkloadCommandBase
         protected readonly SdkFeatureBand _sdkFeatureBand;
         protected readonly ReleaseVersion _targetSdkVersion;
         protected readonly string _fromRollbackDefinition;
-        protected string _workloadSetVersion;
+        protected string _workloadSetVersionFromCommandLine;
+        protected string _globalJsonPath;
         protected string _workloadSetVersionFromGlobalJson;
         protected readonly PackageSourceLocation _packageSourceLocation;
         protected readonly IWorkloadResolverFactory _workloadResolverFactory;
@@ -44,6 +46,10 @@ internal abstract class InstallingWorkloadCommand : WorkloadCommandBase
         protected IInstaller _workloadInstaller;
         protected IWorkloadManifestUpdater _workloadManifestUpdater;
 
+        protected bool UseRollback => !string.IsNullOrWhiteSpace(_fromRollbackDefinition);
+        protected bool SpecifiedWorkloadSetVersionOnCommandLine => !string.IsNullOrWhiteSpace(_workloadSetVersionFromCommandLine);
+        protected bool SpecifiedWorkloadSetVersionInGlobalJson => !string.IsNullOrWhiteSpace(_workloadSetVersionFromGlobalJson);
+
         public InstallingWorkloadCommand(
             ParseResult parseResult,
             IReporter reporter,
@@ -60,6 +66,8 @@ public InstallingWorkloadCommand(
             _downloadToCacheOption = parseResult.GetValue(InstallingWorkloadCommandParser.DownloadToCacheOption);
 
             _fromRollbackDefinition = parseResult.GetValue(InstallingWorkloadCommandParser.FromRollbackFileOption);
+            _workloadSetVersionFromCommandLine = parseResult.GetValue(InstallingWorkloadCommandParser.WorkloadSetVersionOption);
+
             var configOption = parseResult.GetValue(InstallingWorkloadCommandParser.ConfigOption);
             var sourceOption = parseResult.GetValue(InstallingWorkloadCommandParser.SourceOption);
             _packageSourceLocation = string.IsNullOrEmpty(configOption) && (sourceOption == null || !sourceOption.Any()) ? null :
@@ -91,70 +99,162 @@ public InstallingWorkloadCommand(
 
             _workloadInstallerFromConstructor = workloadInstaller;
             _workloadManifestUpdaterFromConstructor = workloadManifestUpdater;
+
+            _globalJsonPath = SdkDirectoryWorkloadManifestProvider.GetGlobalJsonPath(Environment.CurrentDirectory);
+            _workloadSetVersionFromGlobalJson = SdkDirectoryWorkloadManifestProvider.GlobalJsonReader.GetWorkloadVersionFromGlobalJson(_globalJsonPath);
+
+            if (SpecifiedWorkloadSetVersionInGlobalJson && (SpecifiedWorkloadSetVersionOnCommandLine || UseRollback))
+            {
+                throw new GracefulException(string.Format(Strings.CannotSpecifyVersionOnCommandLineAndInGlobalJson, _globalJsonPath), isUserError: true);
+            }
+
+            if (SpecifiedWorkloadSetVersionOnCommandLine && UseRollback)
+            {
+                throw new GracefulException(string.Format(Update.LocalizableStrings.CannotCombineOptions,
+                    InstallingWorkloadCommandParser.FromRollbackFileOption.Name,
+                    InstallingWorkloadCommandParser.WorkloadSetVersionOption.Name), isUserError: true);
+            }
+
+            //  At this point, at most one of SpecifiedWorkloadSetVersionOnCommandLine, UseRollback, and SpecifiedWorkloadSetVersionInGlobalJson is true
         }
 
         protected static Dictionary<string, string> GetInstallStateContents(IEnumerable<ManifestVersionUpdate> manifestVersionUpdates) =>
             WorkloadSet.FromManifests(
                     manifestVersionUpdates.Select(update => new WorkloadManifestInfo(update.ManifestId.ToString(), update.NewVersion.ToString(), /* We don't actually use the directory here */ string.Empty, update.NewFeatureBand))
                     ).ToDictionaryForJson();
 
-        public static bool ShouldUseWorkloadSetMode(SdkFeatureBand sdkFeatureBand, string dotnetDir)
+        InstallStateContents GetCurrentInstallState()
         {
-            string path = Path.Combine(WorkloadInstallType.GetInstallStateFolder(sdkFeatureBand, dotnetDir), "default.json");
-            var installStateContents = File.Exists(path) ? InstallStateContents.FromString(File.ReadAllText(path)) : new InstallStateContents();
-            return installStateContents.UseWorkloadSets ?? false;
+            return GetCurrentInstallState(_sdkFeatureBand, _dotnetPath);
         }
 
-        protected void ErrorIfGlobalJsonAndCommandLineMismatch(string globaljsonPath)
+        static InstallStateContents GetCurrentInstallState(SdkFeatureBand sdkFeatureBand, string dotnetDir)
         {
-            if (!string.IsNullOrWhiteSpace(_workloadSetVersionFromGlobalJson) && !string.IsNullOrWhiteSpace(_workloadSetVersion) && !_workloadSetVersion.Equals(_workloadSetVersionFromGlobalJson))
-            {
-                throw new Exception(string.Format(Strings.CannotSpecifyVersionOnCommandLineAndInGlobalJson, globaljsonPath));
-            }
+            string path = Path.Combine(WorkloadInstallType.GetInstallStateFolder(sdkFeatureBand, dotnetDir), "default.json");
+            return InstallStateContents.FromPath(path);
         }
 
-        protected bool TryHandleWorkloadUpdateFromVersion(ITransactionContext context, DirectoryPath? offlineCache, out IEnumerable<ManifestVersionUpdate> updates)
+        public static bool ShouldUseWorkloadSetMode(SdkFeatureBand sdkFeatureBand, string dotnetDir)
         {
-            // Ensure workload set mode is set to 'workloadset'
-            // Do not skip checking the mode first, as setting it triggers
-            // an admin authorization popup for MSI-based installs.
-            if (!ShouldUseWorkloadSetMode(_sdkFeatureBand, _dotnetPath))
-            {
-                _workloadInstaller.UpdateInstallMode(_sdkFeatureBand, true);
-            }
-
-            _workloadManifestUpdater.DownloadWorkloadSet(_workloadSetVersionFromGlobalJson ?? _workloadSetVersion, offlineCache);
-            return TryInstallWorkloadSet(context, out updates, throwOnFailure: true);
+            return GetCurrentInstallState(sdkFeatureBand, dotnetDir).UseWorkloadSets ?? false;
         }
 
-        public bool TryInstallWorkloadSet(ITransactionContext context, out IEnumerable<ManifestVersionUpdate> updates, bool throwOnFailure = false)
+        protected void UpdateWorkloadManifests(ITransactionContext context, DirectoryPath? offlineCache)
         {
-            var advertisingPackagePath = Path.Combine(_userProfileDir, "sdk-advertising", _sdkFeatureBand.ToString(), "microsoft.net.workloads");
-            if (File.Exists(Path.Combine(advertisingPackagePath, Constants.workloadSetVersionFileName)))
+            var updateToLatestWorkloadSet = ShouldUseWorkloadSetMode(_sdkFeatureBand, _dotnetPath);
+            if (UseRollback && updateToLatestWorkloadSet)
             {
-                // This file isn't created in tests.
-                var version = File.ReadAllText(Path.Combine(advertisingPackagePath, Constants.workloadSetVersionFileName));
-                PrintWorkloadSetTransition(version);
+                // Rollback files are only for loose manifests. Update the mode to be loose manifests.
+                Reporter.WriteLine(Update.LocalizableStrings.UpdateFromRollbackSwitchesModeToLooseManifests);
+                _workloadInstaller.UpdateInstallMode(_sdkFeatureBand, false);
+                updateToLatestWorkloadSet = false;
             }
-            else if (_workloadInstaller is FileBasedInstaller || _workloadInstaller is NetSdkMsiInstallerClient)
+
+            if (SpecifiedWorkloadSetVersionOnCommandLine)
             {
-                // No workload sets found
-                if (throwOnFailure)
+                updateToLatestWorkloadSet = false;
+
+                //  If a workload set version is specified, then switch to workload set update mode
+                //  Check to make sure the value needs to be changed, as updating triggers a UAC prompt
+                //  for MSI-based installs.
+                if (!ShouldUseWorkloadSetMode(_sdkFeatureBand, _dotnetPath))
                 {
-                    throw new NuGetPackageNotFoundException(string.Format(Update.LocalizableStrings.WorkloadVersionRequestedNotFound, _workloadSetVersionFromGlobalJson ?? _workloadSetVersion));
+                    _workloadInstaller.UpdateInstallMode(_sdkFeatureBand, true);
                 }
-                else
+            }
+
+            string resolvedWorkloadSetVersion = _workloadSetVersionFromGlobalJson ??_workloadSetVersionFromCommandLine;
+            if (string.IsNullOrWhiteSpace(resolvedWorkloadSetVersion) && !UseRollback)
+            {
+                _workloadManifestUpdater.UpdateAdvertisingManifestsAsync(_includePreviews, updateToLatestWorkloadSet, offlineCache).Wait();
+                if (updateToLatestWorkloadSet)
                 {
-                    Reporter.WriteLine(Update.LocalizableStrings.NoWorkloadUpdateFound);
+                    resolvedWorkloadSetVersion = _workloadManifestUpdater.GetAdvertisedWorkloadSetVersion();
                 }
-                updates = null;
-                return false;
             }
 
-            var workloadSetPath = _workloadInstaller.InstallWorkloadSet(context, advertisingPackagePath);
-            var files = Directory.EnumerateFiles(workloadSetPath, "*.workloadset.json");
-            updates = _workloadManifestUpdater.ParseRollbackDefinitionFiles(files);
-            return true;
+            if (updateToLatestWorkloadSet && resolvedWorkloadSetVersion == null)
+            {
+                Reporter.WriteLine(Update.LocalizableStrings.NoWorkloadUpdateFound);
+                return;
+            }
+
+            IEnumerable<ManifestVersionUpdate> manifestsToUpdate;
+            if (resolvedWorkloadSetVersion != null)
+            {
+                manifestsToUpdate = InstallWorkloadSet(context, resolvedWorkloadSetVersion);
+            }
+            else
+            {
+                manifestsToUpdate = UseRollback ? _workloadManifestUpdater.CalculateManifestRollbacks(_fromRollbackDefinition) :
+                                                  _workloadManifestUpdater.CalculateManifestUpdates().Select(m => m.ManifestUpdate);
+            }
+
+            InstallStateContents oldInstallState = GetCurrentInstallState();
+
+            context.Run(
+                action: () =>
+                {
+                    foreach (var manifestUpdate in manifestsToUpdate)
+                    {
+                        _workloadInstaller.InstallWorkloadManifest(manifestUpdate, context, offlineCache);
+                    }
+
+                    if (!SpecifiedWorkloadSetVersionInGlobalJson)
+                    {
+                        if (UseRollback)
+                        {
+                            _workloadInstaller.SaveInstallStateManifestVersions(_sdkFeatureBand, GetInstallStateContents(manifestsToUpdate));
+                        }
+                        else if (SpecifiedWorkloadSetVersionOnCommandLine)
+                        {
+                            _workloadInstaller.AdjustWorkloadSetInInstallState(_sdkFeatureBand, resolvedWorkloadSetVersion);
+                        }
+                        else if (this is WorkloadUpdateCommand)
+                        {
+                            //  For workload updates, if you don't specify a rollback file, or a workload version then we should update to a new version of the manifests or workload set, and
+                            //  should remove the install state that pins to the other version
+                            _workloadInstaller.RemoveManifestsFromInstallState(_sdkFeatureBand);
+                            _workloadInstaller.AdjustWorkloadSetInInstallState(_sdkFeatureBand, null);
+                        }
+                    }
+
+                    _workloadResolver.RefreshWorkloadManifests();
+                },
+                rollback: () =>
+                {
+                    //  Reset install state
+                    var currentInstallState = GetCurrentInstallState();
+                    if (currentInstallState.UseWorkloadSets != oldInstallState.UseWorkloadSets)
+                    {
+                        _workloadInstaller.UpdateInstallMode(_sdkFeatureBand, oldInstallState.UseWorkloadSets);
+                    }
+
+                    if ((currentInstallState.Manifests == null && oldInstallState.Manifests != null) ||
+                        (currentInstallState.Manifests != null && oldInstallState.Manifests == null) ||
+                        (currentInstallState.Manifests != null && oldInstallState.Manifests != null &&
+                         (currentInstallState.Manifests.Count != oldInstallState.Manifests.Count ||
+                         !currentInstallState.Manifests.All(m => oldInstallState.Manifests.TryGetValue(m.Key, out var val) && val.Equals(m.Value)))))
+                    {
+                        _workloadInstaller.SaveInstallStateManifestVersions(_sdkFeatureBand, oldInstallState.Manifests);
+                    }
+
+                    if (currentInstallState.WorkloadVersion != oldInstallState.WorkloadVersion)
+                    {
+                        _workloadInstaller.AdjustWorkloadSetInInstallState(_sdkFeatureBand, oldInstallState.WorkloadVersion);
+                    }
+
+                    //  We will refresh the workload manifests to make sure that the resolver has the updated state after the rollback
+                    _workloadResolver.RefreshWorkloadManifests();
+                });
+        }
+
+        private IEnumerable<ManifestVersionUpdate> InstallWorkloadSet(ITransactionContext context, string workloadSetVersion)
+        {
+            PrintWorkloadSetTransition(workloadSetVersion);
+            var workloadSet = _workloadInstaller.InstallWorkloadSet(context, workloadSetVersion);
+
+            return _workloadManifestUpdater.CalculateManifestUpdatesForWorkloadSet(workloadSet);
         }
 
         private void PrintWorkloadSetTransition(string newVersion)
@@ -258,6 +358,17 @@ protected IEnumerable<WorkloadId> GetInstalledWorkloads(bool fromPreviousSdk)
                 return workloads ?? Enumerable.Empty<WorkloadId>();
             }
         }
+
+        protected IEnumerable<WorkloadId> WriteSDKInstallRecordsForVSWorkloads(IEnumerable<WorkloadId> workloadsWithExistingInstallRecords)
+        {
+#if !DOT_NET_BUILD_FROM_SOURCE
+            if (OperatingSystem.IsWindows())
+            {
+                return VisualStudioWorkloads.WriteSDKInstallRecordsForVSWorkloads(_workloadInstaller, _workloadResolver, workloadsWithExistingInstallRecords, Reporter);
+            }
+#endif
+            return workloadsWithExistingInstallRecords;
+        }
     }
 
     internal static class InstallingWorkloadCommandParser