Fix handling installer shutdown in workload restore command

Author: dsplaisted

src/Cli/dotnet/commands/dotnet-workload/install/WorkloadInstallCommand.cs

@@ -20,6 +20,7 @@ internal class WorkloadInstallCommand : InstallingWorkloadCommand
     {
         private bool _skipManifestUpdate;
         private readonly IReadOnlyCollection<string> _workloadIds;
+        private readonly bool _shouldShutdownInstaller;
 
         public bool IsRunningRestore { get; set; }
 
@@ -45,6 +46,7 @@ public WorkloadInstallCommand(
                                  WorkloadInstallerFactory.GetWorkloadInstaller(resolvedReporter, _sdkFeatureBand,
                                      _workloadResolver, Verbosity, _userProfileDir, VerifySignatures, PackageDownloader, _dotnetPath, TempDirectoryPath,
                                      _packageSourceLocation, RestoreActionConfiguration, elevationRequired: !_printDownloadLinkOnly && string.IsNullOrWhiteSpace(_downloadToCacheOption));
+            _shouldShutdownInstaller = _workloadInstallerFromConstructor != null;
 
             _workloadManifestUpdater = _workloadManifestUpdaterFromConstructor ?? new WorkloadManifestUpdater(resolvedReporter, _workloadResolver, PackageDownloader, _userProfileDir,
                 _workloadInstaller.GetWorkloadInstallationRecordRepository(), _workloadInstaller, _packageSourceLocation, displayManifestUpdates: Verbosity.IsDetailedOrDiagnostic());
@@ -146,14 +148,21 @@ public override int Execute()
                 }
                 catch (Exception e)
                 {
-                    _workloadInstaller.Shutdown();
+                    if (_shouldShutdownInstaller)
+                    {
+                        _workloadInstaller.Shutdown();
+                    }
 
                     // Don't show entire stack trace
                     throw new GracefulException(string.Format(LocalizableStrings.WorkloadInstallationFailed, e.Message), e, isUserError: false);
                 }
             }
 
-            _workloadInstaller.Shutdown();
+            if (_shouldShutdownInstaller)
+            {
+                _workloadInstaller.Shutdown();
+            }
+            
             return _workloadInstaller.ExitCode;
         }
 

src/Cli/dotnet/commands/dotnet-workload/restore/WorkloadRestoreCommand.cs

@@ -51,7 +51,7 @@ public override int Execute()
             recorder.Run(() =>
             {
                 // First update manifests and install a workload set as necessary
-                new WorkloadUpdateCommand(_result, recorder: recorder).Execute();
+                new WorkloadUpdateCommand(_result, recorder: recorder, isRestoring: true).Execute();
 
                 var allProjects = DiscoverAllProjects(Directory.GetCurrentDirectory(), _slnOrProjectArgument).Distinct();
                 List<WorkloadId> allWorkloadId = RunTargetToGetWorkloadIds(allProjects);
@@ -64,6 +64,8 @@ public override int Execute()
                     IsRunningRestore = true
                 }.Execute();
             });
+
+            workloadInstaller.Shutdown();
 
             return 0;
         }

src/Cli/dotnet/commands/dotnet-workload/update/WorkloadUpdateCommand.cs

@@ -23,6 +23,7 @@ internal class WorkloadUpdateCommand : InstallingWorkloadCommand
         private readonly bool _fromPreviousSdk;
         private WorkloadHistoryRecorder _recorder;
         private readonly bool _isRestoring;
+        private readonly bool _shouldShutdownInstaller;
         public WorkloadUpdateCommand(
             ParseResult parseResult,
             IReporter reporter = null,
@@ -49,6 +50,9 @@ public WorkloadUpdateCommand(
                                 _dotnetPath, TempDirectoryPath, packageSourceLocation: _packageSourceLocation, RestoreActionConfiguration,
                                 elevationRequired: !_printDownloadLinkOnly && !_printRollbackDefinitionOnly && string.IsNullOrWhiteSpace(_downloadToCacheOption));
 
+            _shouldShutdownInstaller = _workloadInstallerFromConstructor != null;
+
+
             _workloadManifestUpdater = _workloadManifestUpdaterFromConstructor ?? new WorkloadManifestUpdater(resolvedReporter, _workloadResolver, PackageDownloader, _userProfileDir,
                 _workloadInstaller.GetWorkloadInstallationRecordRepository(), _workloadInstaller, _packageSourceLocation, sdkFeatureBand: _sdkFeatureBand);
             _recorder = recorder;
@@ -132,7 +136,10 @@ public override int Execute()
                 }
             }
 
-            _workloadInstaller.Shutdown();
+            if (_shouldShutdownInstaller)
+            {
+                _workloadInstaller.Shutdown();
+            }
             return _workloadInstaller.ExitCode;
         }