Refactor workload update and install commands

Author: dsplaisted

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/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;
@@ -104,31 +105,107 @@ public static bool ShouldUseWorkloadSetMode(SdkFeatureBand sdkFeatureBand, strin
             return installStateContents.UseWorkloadSets ?? false;
         }
 
-        protected IEnumerable<ManifestVersionUpdate> HandleWorkloadUpdateFromVersion(ITransactionContext context, DirectoryPath? offlineCache)
+        protected void UpdateWorkloads(ITransactionContext context, IEnumerable<WorkloadId> workloadsToUpdateOrInstall, DirectoryPath? offlineCache)
         {
-            // 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))
+            var useRollback = !string.IsNullOrWhiteSpace(_fromRollbackDefinition);
+
+            bool specifiedWorkloadSetVersion = !string.IsNullOrEmpty(_workloadSetVersion);
+            if (specifiedWorkloadSetVersion && useRollback)
             {
-                _workloadInstaller.UpdateInstallMode(_sdkFeatureBand, true);
+                throw new GracefulException(string.Format(Update.LocalizableStrings.CannotCombineOptions,
+                    InstallingWorkloadCommandParser.FromRollbackFileOption.Name,
+                    InstallingWorkloadCommandParser.WorkloadSetVersionOption.Name), isUserError: true);
             }
 
-            _workloadManifestUpdater.DownloadWorkloadSet(_workloadSetVersion, offlineCache);
-            return InstallWorkloadSet(context);
-        }
+            var useWorkloadSets = ShouldUseWorkloadSetMode(_sdkFeatureBand, _dotnetPath);
+            if (useRollback && useWorkloadSets)
+            {
+                // Rollback files are only for loose manifests. Update the mode to be loose manifests.
+                Reporter.WriteLine(Update.LocalizableStrings.UpdateFromRollbackSwitchesModeToLooseManifests);
+                _workloadInstaller.UpdateInstallMode(_sdkFeatureBand, false);
+                useWorkloadSets = false;
+            }
 
-        public IEnumerable<ManifestVersionUpdate> InstallWorkloadSet(ITransactionContext context)
-        {
-            var advertisingPackagePath = Path.Combine(_userProfileDir, "sdk-advertising", _sdkFeatureBand.ToString(), "microsoft.net.workloads");
-            if (File.Exists(Path.Combine(advertisingPackagePath, Constants.workloadSetVersionFileName)))
+            if (specifiedWorkloadSetVersion)
+            {
+                useWorkloadSets = true;
+
+                //  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))
+                {
+                    _workloadInstaller.UpdateInstallMode(_sdkFeatureBand, true);
+                }
+            }
+
+            string resolvedWorkloadSetVersion = _workloadSetVersion;
+            if (string.IsNullOrWhiteSpace(_workloadSetVersion))
+            {
+                _workloadManifestUpdater.UpdateAdvertisingManifestsAsync(_includePreviews, useWorkloadSets, offlineCache).Wait();
+                if (useWorkloadSets)
+                {
+                    resolvedWorkloadSetVersion = _workloadManifestUpdater.GetAdvertisedWorkloadSetVersion();
+                }
+            }
+
+            if (useWorkloadSets && resolvedWorkloadSetVersion == null)
+            {
+                Reporter.WriteLine(Update.LocalizableStrings.NoWorkloadUpdateFound);
+                return;
+            }
+
+            IEnumerable<ManifestVersionUpdate> manifestsToUpdate;
+            if (useWorkloadSets)
+            {
+                manifestsToUpdate = InstallWorkloadSet(context, resolvedWorkloadSetVersion);
+            }
+            else
             {
-                // This file isn't created in tests.
-                PrintWorkloadSetTransition(File.ReadAllText(Path.Combine(advertisingPackagePath, Constants.workloadSetVersionFileName)));
+                manifestsToUpdate = useRollback ? _workloadManifestUpdater.CalculateManifestRollbacks(_fromRollbackDefinition) :
+                                                  _workloadManifestUpdater.CalculateManifestUpdates().Select(m => m.ManifestUpdate);
             }
-            var workloadSetPath = _workloadInstaller.InstallWorkloadSet(context, advertisingPackagePath);
-            var files = Directory.EnumerateFiles(workloadSetPath, "*.workloadset.json");
-            return _workloadManifestUpdater.ParseRollbackDefinitionFiles(files);
+
+            context.Run(
+                action: () =>
+                {
+                    foreach (var manifestUpdate in manifestsToUpdate)
+                    {
+                        _workloadInstaller.InstallWorkloadManifest(manifestUpdate, context, offlineCache, useRollback);
+                    }
+
+                    if (useRollback)
+                    {
+                        _workloadInstaller.SaveInstallStateManifestVersions(_sdkFeatureBand, GetInstallStateContents(manifestsToUpdate));
+                    }
+                    else if (this is WorkloadUpdateCommand)
+                    {
+                        //  For workload updates, if you don't specify a rollback file, 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
+                        //  TODO: Do we need to do something similar if the workload set version is pinned in the install state?
+                        _workloadInstaller.RemoveManifestsFromInstallState(_sdkFeatureBand);
+                    }
+
+                    _workloadInstaller.AdjustWorkloadSetInInstallState(_sdkFeatureBand, string.IsNullOrWhiteSpace(_workloadSetVersion) ? null : _workloadSetVersion);
+
+                    _workloadResolver.RefreshWorkloadManifests();
+
+                    _workloadInstaller.InstallWorkloads(workloadsToUpdateOrInstall, _sdkFeatureBand, context, offlineCache);
+                },
+                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();
+                });
+        }
+
+        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)
@@ -240,6 +317,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