workload repair can recover from corrupt workload sets

src/Cli/dotnet/Commands/CliCommandStrings.resx

@@ -1,17 +1,17 @@
 <?xml version="1.0" encoding="utf-8"?>
 <root>
-  <!-- 
-    Microsoft ResX Schema 
-    
+  <!--
+    Microsoft ResX Schema
+
     Version 2.0
-    
-    The primary goals of this format is to allow a simple XML format 
-    that is mostly human readable. The generation and parsing of the 
-    various data types are done through the TypeConverter classes 
+
+    The primary goals of this format is to allow a simple XML format
+    that is mostly human readable. The generation and parsing of the
+    various data types are done through the TypeConverter classes
     associated with the data types.
-    
+
     Example:
-    
+
     ... ado.net/XML headers & schema ...
     <resheader name="resmimetype">text/microsoft-resx</resheader>
     <resheader name="version">2.0</resheader>
@@ -26,36 +26,36 @@
         <value>[base64 mime encoded string representing a byte array form of the .NET Framework object]</value>
         <comment>This is a comment</comment>
     </data>
-                
-    There are any number of "resheader" rows that contain simple 
+
+    There are any number of "resheader" rows that contain simple
     name/value pairs.
-    
-    Each data row contains a name, and value. The row also contains a 
-    type or mimetype. Type corresponds to a .NET class that support 
-    text/value conversion through the TypeConverter architecture. 
-    Classes that don't support this are serialized and stored with the 
+
+    Each data row contains a name, and value. The row also contains a
+    type or mimetype. Type corresponds to a .NET class that support
+    text/value conversion through the TypeConverter architecture.
+    Classes that don't support this are serialized and stored with the
     mimetype set.
-    
-    The mimetype is used for serialized objects, and tells the 
-    ResXResourceReader how to depersist the object. This is currently not 
+
+    The mimetype is used for serialized objects, and tells the
+    ResXResourceReader how to depersist the object. This is currently not
     extensible. For a given mimetype the value must be set accordingly:
-    
-    Note - application/x-microsoft.net.object.binary.base64 is the format 
-    that the ResXResourceWriter will generate, however the reader can 
+
+    Note - application/x-microsoft.net.object.binary.base64 is the format
+    that the ResXResourceWriter will generate, however the reader can
     read any of the formats listed below.
-    
+
     mimetype: application/x-microsoft.net.object.binary.base64
-    value   : The object must be serialized with 
+    value   : The object must be serialized with
             : System.Runtime.Serialization.Formatters.Binary.BinaryFormatter
             : and then encoded with base64 encoding.
-    
+
     mimetype: application/x-microsoft.net.object.soap.base64
-    value   : The object must be serialized with 
+    value   : The object must be serialized with
             : System.Runtime.Serialization.Formatters.Soap.SoapFormatter
             : and then encoded with base64 encoding.
 
     mimetype: application/x-microsoft.net.object.bytearray.base64
-    value   : The object must be serialized into a byte array 
+    value   : The object must be serialized into a byte array
             : using a System.ComponentModel.TypeConverter
             : and then encoded with base64 encoding.
     -->
@@ -2724,4 +2724,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

@@ -69,7 +69,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())
                 {
@@ -79,7 +80,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));
 
@@ -105,4 +106,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,121 @@
+// 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?.HasMissingManifests(workloadSet) ?? true)
+        {
+            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);
+    }
+}