workload repair can recover from corrupt workload sets (#52434)

src/Cli/dotnet/Commands/CliCommandStrings.resx

@@ -2651,4 +2651,8 @@ Proceed?</value>
   <data name="DurationColon" xml:space="preserve">
     <value>duration:</value>
   </data>
+  <data name="WorkloadSetHasMissingManifests" xml:space="preserve">
+    <value>Workload set version {0} has missing manifests likely removed by package management. Run "dotnet workload repair" to fix this.</value>
+    <comment>{0} is the workload set version. {Locked="dotnet workload repair"}</comment>
+  </data>
 </root>

src/Cli/dotnet/Commands/Workload/Install/FileBasedInstaller.cs

@@ -5,6 +5,7 @@
 
 using System.Collections.Concurrent;
 using System.Text.Json;
+using Microsoft.DotNet.Cli.Commands.Workload;
 using Microsoft.DotNet.Cli.Commands.Workload.Config;
 using Microsoft.DotNet.Cli.Commands.Workload.Install.WorkloadInstallRecords;
 using Microsoft.DotNet.Cli.Extensions;
@@ -642,7 +643,7 @@ public IEnumerable<WorkloadHistoryRecord> GetWorkloadHistoryRecords(string sdkFe
     public void Shutdown()
     {
         // Perform any additional cleanup here that's intended to run at the end of the command, regardless
-        // of success or failure. For file based installs, there shouldn't be any additional work to 
+        // of success or failure. For file based installs, there shouldn't be any additional work to
         // perform.
     }
 

src/Cli/dotnet/Commands/Workload/Install/WorkloadInstallerFactory.cs

@@ -3,6 +3,7 @@
 
 #nullable disable
 
+using Microsoft.DotNet.Cli.Commands.Workload;
 using Microsoft.DotNet.Cli.NuGetPackageDownloader;
 using Microsoft.DotNet.Cli.Utils;
 using Microsoft.DotNet.Configurer;
@@ -48,7 +49,7 @@ public static IInstaller GetWorkloadInstaller(
 
         userProfileDir ??= CliFolderPathCalculator.DotnetUserProfileFolderPath;
 
-        return new FileBasedInstaller(
+        var installer = new FileBasedInstaller(
             reporter,
             sdkFeatureBand,
             workloadResolver,
@@ -59,6 +60,25 @@ public static IInstaller GetWorkloadInstaller(
             verbosity: verbosity,
             packageSourceLocation: packageSourceLocation,
             restoreActionConfig: restoreActionConfig);
+
+        // Attach corruption repairer to recover from corrupt workload sets
+        if (nugetPackageDownloader is not null &&
+            workloadResolver?.GetWorkloadManifestProvider() is SdkDirectoryWorkloadManifestProvider sdkProvider &&
+            sdkProvider.CorruptionRepairer is null)
+        {
+            sdkProvider.CorruptionRepairer = new WorkloadManifestCorruptionRepairer(
+                reporter,
+                installer,
+                workloadResolver,
+                sdkFeatureBand,
+                dotnetDir,
+                userProfileDir,
+                nugetPackageDownloader,
+                packageSourceLocation,
+                verbosity);
+        }
+
+        return installer;
     }
 
     private static bool CanWriteToDotnetRoot(string dotnetDir = null)

src/Cli/dotnet/Commands/Workload/Repair/WorkloadRepairCommand.cs

@@ -70,7 +70,8 @@ public override int Execute()
             {
                 Reporter.WriteLine();
 
-                var workloadIds = _workloadInstaller.GetWorkloadInstallationRecordRepository().GetInstalledWorkloads(new SdkFeatureBand(_sdkVersion));
+                var sdkFeatureBand = new SdkFeatureBand(_sdkVersion);
+                var workloadIds = _workloadInstaller.GetWorkloadInstallationRecordRepository().GetInstalledWorkloads(sdkFeatureBand);
 
                 if (!workloadIds.Any())
                 {
@@ -80,7 +81,7 @@ public override int Execute()
 
                 Reporter.WriteLine(string.Format(CliCommandStrings.RepairingWorkloads, string.Join(" ", workloadIds)));
 
-                ReinstallWorkloadsBasedOnCurrentManifests(workloadIds, new SdkFeatureBand(_sdkVersion));
+                ReinstallWorkloadsBasedOnCurrentManifests(workloadIds, sdkFeatureBand);
 
                 WorkloadInstallCommand.TryRunGarbageCollection(_workloadInstaller, Reporter, Verbosity, workloadSetVersion => _workloadResolverFactory.CreateForWorkloadSet(_dotnetPath, _sdkVersion.ToString(), _userProfileDir, workloadSetVersion));
 
@@ -106,4 +107,5 @@ private void ReinstallWorkloadsBasedOnCurrentManifests(IEnumerable<WorkloadId> w
     {
         _workloadInstaller.RepairWorkloads(workloadIds, sdkFeatureBand);
     }
+
 }

src/Cli/dotnet/Commands/Workload/WorkloadHistoryRecorder.cs

@@ -54,6 +54,10 @@ public void Run(Action workloadAction)
     private WorkloadHistoryState GetWorkloadState()
     {
         var resolver = _workloadResolverFunc();
+        if (resolver.GetWorkloadManifestProvider() is SdkDirectoryWorkloadManifestProvider sdkProvider)
+        {
+            sdkProvider.CorruptionFailureMode = ManifestCorruptionFailureMode.Ignore;
+        }
         var currentWorkloadVersion = resolver.GetWorkloadVersion().Version;
         return new WorkloadHistoryState()
         {

src/Cli/dotnet/Commands/Workload/WorkloadManifestCorruptionRepairer.cs

@@ -0,0 +1,119 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System.Diagnostics.CodeAnalysis;
+using Microsoft.DotNet.Cli;
+using Microsoft.DotNet.Cli.Commands.Workload.Install;
+using Microsoft.DotNet.Cli.NuGetPackageDownloader;
+using Microsoft.DotNet.Cli.Utils;
+using Microsoft.NET.Sdk.WorkloadManifestReader;
+
+namespace Microsoft.DotNet.Cli.Commands.Workload;
+
+internal sealed class WorkloadManifestCorruptionRepairer : IWorkloadManifestCorruptionRepairer
+{
+    private readonly IReporter _reporter;
+    private readonly IInstaller _workloadInstaller;
+    private readonly IWorkloadResolver _workloadResolver;
+    private readonly SdkFeatureBand _sdkFeatureBand;
+    private readonly string _dotnetPath;
+    private readonly string _userProfileDir;
+    private readonly INuGetPackageDownloader? _packageDownloader;
+    private readonly PackageSourceLocation? _packageSourceLocation;
+    private readonly VerbosityOptions _verbosity;
+
+    private bool _checked;
+
+    public WorkloadManifestCorruptionRepairer(
+        IReporter reporter,
+        IInstaller workloadInstaller,
+        IWorkloadResolver workloadResolver,
+        SdkFeatureBand sdkFeatureBand,
+        string dotnetPath,
+        string userProfileDir,
+        INuGetPackageDownloader? packageDownloader,
+        PackageSourceLocation? packageSourceLocation,
+        VerbosityOptions verbosity)
+    {
+        _reporter = reporter ?? NullReporter.Instance;
+        _workloadInstaller = workloadInstaller;
+        _workloadResolver = workloadResolver;
+        _sdkFeatureBand = sdkFeatureBand;
+        _dotnetPath = dotnetPath;
+        _userProfileDir = userProfileDir;
+        _packageDownloader = packageDownloader;
+        _packageSourceLocation = packageSourceLocation;
+        _verbosity = verbosity;
+    }
+
+    public void EnsureManifestsHealthy(ManifestCorruptionFailureMode failureMode)
+    {
+        if (_checked)
+        {
+            return;
+        }
+
+        _checked = true;
+
+        if (failureMode == ManifestCorruptionFailureMode.Ignore)
+        {
+            return;
+        }
+
+        // Get the workload set directly from the provider - it was already resolved during construction
+        // and doesn't require reading the install state file again
+        var provider = _workloadResolver.GetWorkloadManifestProvider() as SdkDirectoryWorkloadManifestProvider;
+        var workloadSet = provider?.ResolvedWorkloadSet;
+
+        if (workloadSet is null)
+        {
+            // No workload set is being used
+            return;
+        }
+
+        if (provider == null || !provider.HasMissingManifests(workloadSet))
+        {
+            return;
+        }
+
+        if (failureMode == ManifestCorruptionFailureMode.Throw)
+        {
+            throw new InvalidOperationException(string.Format(CliCommandStrings.WorkloadSetHasMissingManifests, workloadSet.Version));
+        }
+
+        _reporter.WriteLine($"Repairing workload set {workloadSet.Version}...");
+        CliTransaction.RunNew(context => RepairCorruptWorkloadSet(context, workloadSet));
+    }
+
+    private void RepairCorruptWorkloadSet(ITransactionContext context, WorkloadSet workloadSet)
+    {
+        var manifestUpdates = CreateManifestUpdatesFromWorkloadSet(workloadSet);
+
+        foreach (var manifestUpdate in manifestUpdates)
+        {
+            _workloadInstaller.InstallWorkloadManifest(manifestUpdate, context);
+        }
+
+    }
+
+    [MemberNotNull(nameof(_packageDownloader))]
+    private IEnumerable<ManifestVersionUpdate> CreateManifestUpdatesFromWorkloadSet(WorkloadSet workloadSet)
+    {
+        if (_packageDownloader is null)
+        {
+            throw new InvalidOperationException("Package downloader is required to repair workload manifests.");
+        }
+
+        var manifestUpdater = new WorkloadManifestUpdater(
+            _reporter,
+            _workloadResolver,
+            _packageDownloader,
+            _userProfileDir,
+            _workloadInstaller.GetWorkloadInstallationRecordRepository(),
+            _workloadInstaller,
+            _packageSourceLocation,
+            displayManifestUpdates: _verbosity >= VerbosityOptions.detailed);
+
+        return manifestUpdater.CalculateManifestUpdatesForWorkloadSet(workloadSet);
+    }
+}