Exception throwing

Author: CyrusNajmabadi

src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs

@@ -11,6 +11,7 @@
 using System.IO;
 using System.Linq;
 using System.Reflection;
+using System.Runtime.CompilerServices;
 using System.Text.Json;
 using System.Threading;
 using System.Threading.Tasks;
@@ -41,6 +42,8 @@ internal sealed class ExtensionMessageHandlerService(
     IExtensionMessageHandlerFactory customMessageHandlerFactory)
     : IExtensionMessageHandlerService
 {
+    private static readonly ConditionalWeakTable<IExtensionMessageHandlerWrapper, IExtensionMessageHandlerWrapper> s_disabledExtensionHandlers = new();
+
     private readonly SolutionServices _solutionServices = solutionServices;
     private readonly IExtensionMessageHandlerFactory _customMessageHandlerFactory = customMessageHandlerFactory;
 
@@ -226,19 +229,28 @@ private async ValueTask<string> HandleExtensionMessageInCurrentProcessAsync<TArg
             throw new InvalidOperationException($"Multiple handlers found for message {messageName}.");
 
         var handler = handlers[0];
+        if (s_disabledExtensionHandlers.TryGetValue(handler, out _))
+            throw new InvalidOperationException($"Handler was disabled due to previous exception.");
 
         try
         {
             var message = JsonSerializer.Deserialize(jsonMessage, handler.MessageType);
             var result = await handler.ExecuteAsync(message, executeArgument, cancellationToken).ConfigureAwait(false);
             return JsonSerializer.Serialize(result, handler.ResponseType);
         }
-        catch
+        catch (Exception ex) when (DisableHandlerAndPropagate(ex))
         {
-            // Any exception thrown in this method is left to bubble up to the extension.
-            // But we unregister all handlers from that assembly to minimize the impact of a bad extension.
-            await UnregisterExtensionAsync(assemblyFilePath: handler.ExtensionIdentifier, cancellationToken).ConfigureAwait(false);
-            throw;
+            throw ExceptionUtilities.Unreachable();
+        }
+
+        bool DisableHandlerAndPropagate(Exception ex)
+        {
+            FatalError.ReportNonFatalError(ex, ErrorSeverity.Critical);
+
+            // Any exception thrown in this method is left to bubble up to the extension. But we unregister this handler
+            // from that assembly to minimize the impact.
+            s_disabledExtensionHandlers.TryAdd(handler, handler);
+            return false;
         }
     }
 

src/Features/Core/Portable/Extensions/IExtensionMessageHandlerWrapper.cs

@@ -8,13 +8,17 @@
 
 namespace Microsoft.CodeAnalysis.Extensions;
 
+internal interface IExtensionMessageHandlerWrapper
+{
+}
+
 /// <summary>
 /// Wrapper for an <c>IExtensionWorkspaceMessageHandler</c> or <c>IExtensionDocumentMessageHandler</c>
 /// as returned by <see cref="IExtensionMessageHandlerFactory"/>. 
 /// </summary>
 /// <typeparam name="TArgument">The type of object received as parameter by the extension message
 /// handler.</typeparam>
-internal interface IExtensionMessageHandlerWrapper<TArgument>
+internal interface IExtensionMessageHandlerWrapper<TArgument> : IExtensionMessageHandlerWrapper
 {
     /// <summary>
     /// The type of object received as parameter by the extension message handler.