Merge pull request #7010 from ryanbrandenburg/CLaSPReaction

Fix shutdown issues from CLaSP

eng/Versions.props

@@ -81,7 +81,7 @@
     <RoslynPackageVersion>4.4.0-2.22424.2</RoslynPackageVersion>
     <VisualStudioLanguageServerProtocolVersion>17.4.1008-preview</VisualStudioLanguageServerProtocolVersion>
     <MicrosoftNetCompilersToolsetVersion>4.4.0-1.final</MicrosoftNetCompilersToolsetVersion>
-    <MicrosoftCommonLanguageServerProtocolFrameworkPackageVersion>4.5.0-1.22513.3</MicrosoftCommonLanguageServerProtocolFrameworkPackageVersion>
+    <MicrosoftCommonLanguageServerProtocolFrameworkPackageVersion>4.5.0-1.22519.14</MicrosoftCommonLanguageServerProtocolFrameworkPackageVersion>
   </PropertyGroup>
   <PropertyGroup Label="Manual">
     <!-- dotnet/runtime packages -->

src/Razor/benchmarks/Microsoft.AspNetCore.Razor.Microbenchmarks/LanguageServer/RazorCSharpFormattingBenchmark.cs

@@ -130,13 +130,14 @@ namespace Microsoft.AspNetCore.Razor.Microbenchmarks.LanguageServer
         }
 
         [GlobalCleanup]
-        public Task CleanupServerAsync()
+        public async Task CleanupServerAsync()
         {
             File.Delete(_filePath);
 
-            RazorLanguageServer.Dispose();
+            var innerServer = RazorLanguageServer.GetInnerLanguageServerForTesting();
 
-            return Task.CompletedTask;
+            await innerServer.ShutdownAsync();
+            await innerServer.ExitAsync();
         }
 
         private void EnsureServicesInitialized()

src/Razor/benchmarks/Microsoft.AspNetCore.Razor.Microbenchmarks/LanguageServer/RazorSemanticTokensBenchmark.cs

@@ -101,10 +101,11 @@ namespace Microsoft.AspNetCore.Razor.Microbenchmarks.LanguageServer
         }
 
         [GlobalCleanup]
-        public Task CleanupServerAsync()
+        public async Task CleanupServerAsync()
         {
-            RazorLanguageServer.Dispose();
-            return Task.CompletedTask;
+            var innerServer = RazorLanguageServer.GetInnerLanguageServerForTesting();
+            await innerServer.ShutdownAsync();
+            await innerServer.ExitAsync();
         }
 
         protected internal override void Builder(IServiceCollection collection)

src/Razor/benchmarks/Microsoft.AspNetCore.Razor.Microbenchmarks/LanguageServer/RazorSemanticTokensScrollingBenchmark.cs

@@ -113,11 +113,12 @@ namespace Microsoft.AspNetCore.Razor.Microbenchmarks.LanguageServer
         }
 
         [GlobalCleanup]
-        public Task CleanupServerAsync()
+        public async Task CleanupServerAsync()
         {
-            RazorLanguageServer.Dispose();
+            var innerServer = RazorLanguageServer.GetInnerLanguageServerForTesting();
 
-            return Task.CompletedTask;
+            await innerServer.ShutdownAsync();
+            await innerServer.ExitAsync();
         }
 
         protected internal override void Builder(IServiceCollection collection)

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer.Common/DefaultRemoteTextLoaderFactory.cs

@@ -50,7 +50,7 @@ namespace Microsoft.AspNetCore.Razor.LanguageServer.Common
                 _filePath = filePath;
             }
 
-            public override Task<TextAndVersion> LoadTextAndVersionAsync(Workspace workspace, DocumentId documentId, CancellationToken cancellationToken)
+            public override Task<TextAndVersion> LoadTextAndVersionAsync(Workspace? workspace, DocumentId? documentId, CancellationToken cancellationToken)
             {
 
                 TextAndVersion textAndVersion;

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/DocumentSnapshotTextLoader.cs

@@ -23,7 +23,7 @@ namespace Microsoft.AspNetCore.Razor.LanguageServer
             _documentSnapshot = documentSnapshot;
         }
 
-        public override async Task<TextAndVersion> LoadTextAndVersionAsync(Workspace workspace, DocumentId documentId, CancellationToken cancellationToken)
+        public override async Task<TextAndVersion> LoadTextAndVersionAsync(Workspace? workspace, DocumentId? documentId, CancellationToken cancellationToken)
         {
             var sourceText = await _documentSnapshot.GetTextAsync();
             var textAndVersion = TextAndVersion.Create(sourceText, VersionStamp.Default);

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Extensions/IServiceCollectionExtensions.cs

@@ -32,7 +32,6 @@ internal static class IServiceCollectionExtensions
     {
         services.AddHandler<RazorInitializeEndpoint>();
         services.AddHandler<RazorInitializedEndpoint>();
-        services.AddHandler<ShutdownHandler<RazorRequestContext>>();
 
         var razorLifeCycleManager = new RazorLifeCycleManager(razorLanguageServer);
         services.AddSingleton<ILifeCycleManager>(razorLifeCycleManager);

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/LspServices.cs

@@ -12,6 +12,7 @@ namespace Microsoft.AspNetCore.Razor.LanguageServer;
 internal class LspServices : ILspServices
 {
     private readonly IServiceProvider _serviceProvider;
+    public bool IsDisposed = false;
 
     public LspServices(IServiceCollection serviceCollection)
     {
@@ -57,6 +58,7 @@ internal class LspServices : ILspServices
         if (_serviceProvider is IDisposable disposable)
         {
             disposable.Dispose();
+            IsDisposed = true;
         }
     }
 }

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/ProjectSystem/DefaultRazorProjectService.cs

@@ -497,8 +497,8 @@ namespace Microsoft.AspNetCore.Razor.LanguageServer.ProjectSystem
                 _fromDocument = fromDocument ?? throw new ArgumentNullException(nameof(fromDocument));
             }
             public override async Task<TextAndVersion> LoadTextAndVersionAsync(
-               Workspace workspace,
-               DocumentId documentId,
+               Workspace? workspace,
+               DocumentId? documentId,
                CancellationToken cancellationToken)
             {
                 var sourceText = await _fromDocument.GetTextAsync();

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/RazorLanguageServerWrapper.cs

@@ -7,6 +7,7 @@ using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor.LanguageServer.Common.Extensions;
 using Microsoft.CodeAnalysis.Razor;
 using Microsoft.CodeAnalysis.Razor.Workspaces;
+using Microsoft.CommonLanguageServerProtocol.Framework;
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.VisualStudio.LanguageServer.Protocol;
 using Newtonsoft.Json.Serialization;
@@ -68,9 +69,24 @@ internal sealed class RazorLanguageServerWrapper : IDisposable
 
     public Task WaitForExitAsync()
     {
-        var lifeCycleManager = GetRequiredService<RazorLifeCycleManager>();
-
-        return lifeCycleManager.WaitForExit;
+        var lspServices = _innerServer.GetLspServices();
+        if (lspServices is LspServices razorServices)
+        {
+            // If the LSP Server is already disposed it means the server has already exited.
+            if (razorServices.IsDisposed)
+            {
+                return Task.CompletedTask;
+            }
+            else
+            {
+                var lifeCycleManager = razorServices.GetRequiredService<RazorLifeCycleManager>();
+                return lifeCycleManager.WaitForExit;
+            }
+        }
+        else
+        {
+            throw new NotImplementedException($"LspServices should always be of type {nameof(LspServices)}.");
+        }
     }
 
     internal T GetRequiredService<T>() where T : notnull

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/RazorLifeCycleManager.cs

@@ -16,17 +16,18 @@ namespace Microsoft.AspNetCore.Razor.LanguageServer
             _languageServer = languageServer;
         }
 
-        public async Task ExitAsync()
+        public Task ExitAsync()
         {
-            await _languageServer.ExitAsync();
             var services = _languageServer.GetLspServices();
             services.Dispose();
             _tcs.TrySetResult(0);
+
+            return Task.CompletedTask;
         }
 
-        public async Task ShutdownAsync(string message = "Shutting down")
+        public Task ShutdownAsync(string message = "Shutting down")
         {
-            await _languageServer.ShutdownAsync(message);
+            return Task.CompletedTask;
         }
 
         public Task WaitForExit => _tcs.Task;

src/Razor/src/Microsoft.VisualStudio.LanguageServerClient.Razor/CSharpVirtualDocumentPublisher.cs

@@ -158,7 +158,7 @@ namespace Microsoft.VisualStudio.LanguageServerClient.Razor
                     _filePath = filePath;
                 }
 
-                public override Task<TextAndVersion> LoadTextAndVersionAsync(CodeAnalysisWorkspace workspace, DocumentId documentId, CancellationToken cancellationToken)
+                public override Task<TextAndVersion> LoadTextAndVersionAsync(CodeAnalysisWorkspace? workspace, DocumentId? documentId, CancellationToken cancellationToken)
                 {
                     return Task.FromResult(TextAndVersion.Create(_sourceText, VersionStamp.Default, _filePath));
                 }

src/Razor/src/Microsoft.VisualStudio.LanguageServerClient.Razor/RazorLanguageServerClient.cs

@@ -40,9 +40,7 @@ namespace Microsoft.VisualStudio.LanguageServerClient.Razor
         private readonly RazorLanguageServerLogHubLoggerProviderFactory _logHubLoggerProviderFactory;
         private readonly LanguageServerFeatureOptions _languageServerFeatureOptions;
         private readonly VisualStudioHostServicesProvider? _vsHostWorkspaceServicesProvider;
-        private readonly object _shutdownLock;
         private RazorLanguageServerWrapper? _server;
-        private IDisposable? _serverShutdownDisposable;
         private LogHubLoggerProvider? _loggerProvider;
         private readonly ProjectSnapshotManagerDispatcher _projectSnapshotManagerDispatcher;
 
@@ -108,7 +106,6 @@ namespace Microsoft.VisualStudio.LanguageServerClient.Razor
             _logHubLoggerProviderFactory = logHubLoggerProviderFactory;
             _languageServerFeatureOptions = languageServerFeatureOptions;
             _vsHostWorkspaceServicesProvider = vsHostWorkspaceServicesProvider;
-            _shutdownLock = new object();
             _projectSnapshotManagerDispatcher = projectSnapshotManagerDispatcher;
         }
 
@@ -138,7 +135,7 @@ namespace Microsoft.VisualStudio.LanguageServerClient.Razor
 
             var (clientStream, serverStream) = FullDuplexStream.CreatePair();
 
-            await EnsureCleanedUpServerAsync(token).ConfigureAwait(false);
+            await EnsureCleanedUpServerAsync().ConfigureAwait(false);
 
             var traceLevel = GetVerbosity();
 
@@ -176,46 +173,19 @@ namespace Microsoft.VisualStudio.LanguageServerClient.Razor
             return result;
         }
 
-        private async Task EnsureCleanedUpServerAsync(CancellationToken token)
+        private async Task EnsureCleanedUpServerAsync()
         {
-            const int WaitForShutdownAttempts = 10;
-
             if (_server is null)
             {
                 // Server was already cleaned up
                 return;
             }
 
-            var attempts = 0;
-            while (_server is not null && ++attempts < WaitForShutdownAttempts)
-            {
-                // Server failed to shutdown, lets wait a little bit and check again.
-                await Task.Delay(100, token).ConfigureAwait(false);
-            }
-
             if (_server is not null)
             {
-                // Server still hasn't shutdown, attempt an ungraceful shutdown.
-                _server.Dispose();
-
-                ServerShutdown();
-            }
-        }
-
-        private void ServerShutdown()
-        {
-            lock (_shutdownLock)
-            {
-                if (_server is null)
-                {
-                    // Already shutdown
-                    return;
-                }
-
                 _projectConfigurationFilePathStore.Changed -= ProjectConfigurationFilePathStore_Changed;
-                _serverShutdownDisposable?.Dispose();
-                _serverShutdownDisposable = null;
-                _server = null;
+                // Server still hasn't shutdown, wait for it to shutdown
+                await _server.WaitForExitAsync().ConfigureAwait(false);
             }
         }
 

src/Razor/test/Microsoft.AspNetCore.Razor.LanguageServer.Test/RazorLanguageServerTest.cs

@@ -36,8 +36,7 @@ public class RazorLanguageServerTest : TestBase
         // We turn this into a Set to handle cases like Completion where we have two handlers, only one of which will be registered
         // CLaSP will throw if two handlers register for the same method, so if THAT doesn't hold it's a CLaSP bug, not a Razor bug.
         var typeMethods = handlerTypes.Select(t => GetMethodFromType(t)).ToHashSet();
-        // The shutdown handler is outside of our assembly.
-        typeMethods.Add("shutdown");
+
         if (registeredMethods.Length != typeMethods.Count)
         {
             var unregisteredHandlers = typeMethods.Where(t => !registeredMethods.Any(m => m.MethodName == t));