WIP: Implement a safe wrapper around ILanguageServerAdapter to ensure messages are only sent after initialization

src/PowerShellEditorServices/Extensions/Api/EditorContextService.cs

@@ -5,6 +5,7 @@
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.PowerShell.EditorServices.Handlers;
+using Microsoft.PowerShell.EditorServices.Server;
 using OmniSharp.Extensions.LanguageServer.Protocol.Server;
 
 namespace Microsoft.PowerShell.EditorServices.Extensions.Services
@@ -82,10 +83,10 @@ public interface IEditorContextService
 
     internal class EditorContextService : IEditorContextService
     {
-        private readonly ILanguageServerFacade _languageServer;
+        private readonly ISafeLanguageServer _languageServer;
 
         internal EditorContextService(
-            ILanguageServerFacade languageServer)
+            ISafeLanguageServer languageServer)
         {
             _languageServer = languageServer;
         }

src/PowerShellEditorServices/Extensions/Api/EditorExtensionServiceProvider.cs

@@ -5,6 +5,7 @@
 using System.Linq.Expressions;
 using System.Reflection;
 using Microsoft.Extensions.DependencyInjection;
+using Microsoft.PowerShell.EditorServices.Server;
 using Microsoft.PowerShell.EditorServices.Services;
 using Microsoft.PowerShell.EditorServices.Utility;
 using OmniSharp.Extensions.LanguageServer.Protocol.Server;
@@ -41,12 +42,13 @@ public class EditorExtensionServiceProvider
         internal EditorExtensionServiceProvider(IServiceProvider serviceProvider)
         {
             _serviceProvider = serviceProvider;
-            LanguageServer = new LanguageServerService(_serviceProvider.GetService<ILanguageServerFacade>());
+            var languageServer = _serviceProvider.GetService<ISafeLanguageServer>();
+            LanguageServer = new LanguageServerService(languageServer);
             //DocumentSymbols = new DocumentSymbolService(_serviceProvider.GetService<SymbolsService>());
             ExtensionCommands = new ExtensionCommandService(_serviceProvider.GetService<ExtensionService>());
             Workspace = new WorkspaceService(_serviceProvider.GetService<InternalServices.WorkspaceService>());
-            EditorContext = new EditorContextService(_serviceProvider.GetService<ILanguageServerFacade>());
-            EditorUI = new EditorUIService(_serviceProvider.GetService<ILanguageServerFacade>());
+            EditorContext = new EditorContextService(languageServer);
+            EditorUI = new EditorUIService(languageServer);
         }
 
         /// <summary>

src/PowerShellEditorServices/Extensions/Api/EditorUIService.cs

@@ -6,6 +6,7 @@
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.PowerShell.EditorServices.Server;
 using Microsoft.PowerShell.EditorServices.Services.PowerShellContext;
 using OmniSharp.Extensions.LanguageServer.Protocol.Server;
 
@@ -101,9 +102,9 @@ internal class EditorUIService : IEditorUIService
     {
         private static string[] s_choiceResponseLabelSeparators = new[] { ", " };
 
-        private readonly ILanguageServerFacade _languageServer;
+        private readonly ISafeLanguageServer _languageServer;
 
-        public EditorUIService(ILanguageServerFacade languageServer)
+        public EditorUIService(ISafeLanguageServer languageServer)
         {
             _languageServer = languageServer;
         }

src/PowerShellEditorServices/Extensions/Api/LanguageServerService.cs

@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 using MediatR;
+using Microsoft.PowerShell.EditorServices.Server;
 using OmniSharp.Extensions.LanguageServer.Protocol.Server;
 using System.Threading;
 using System.Threading.Tasks;
@@ -64,9 +65,9 @@ public interface ILanguageServerService
 
     internal class LanguageServerService : ILanguageServerService
     {
-        private readonly ILanguageServerFacade _languageServer;
+        private readonly ISafeLanguageServer _languageServer;
 
-        internal LanguageServerService(ILanguageServerFacade languageServer)
+        internal LanguageServerService(ISafeLanguageServer languageServer)
         {
             _languageServer = languageServer;
         }

src/PowerShellEditorServices/Hosting/EditorServicesServerFactory.cs

@@ -123,6 +123,7 @@ public PsesDebugServer CreateDebugServerForTempSession(Stream inputStream, Strea
                 .AddSingleton<ILanguageServerFacade>(provider => null)
                 .AddPsesLanguageServices(hostStartupInfo)
                 // For a Temp session, there is no LanguageServer so just set it to null
+                // TODO: Why are we doing this twice?
                 .AddSingleton(
                     typeof(ILanguageServerFacade),
                     _ => null)

src/PowerShellEditorServices/Server/PsesLanguageServer.cs

@@ -93,7 +93,7 @@ public async Task StartAsync()
                     .WithHandler<PsesSemanticTokensHandler>()
                     .OnInitialize(
                         // TODO: Either fix or ignore "method lacks 'await'" warning.
-                        async (languageServer, request, cancellationToken) =>
+                        (languageServer, request, cancellationToken) =>
                         {
                             var serviceProvider = languageServer.Services;
                             var workspaceService = serviceProvider.GetService<WorkspaceService>();
@@ -113,6 +113,12 @@ public async Task StartAsync()
                                     break;
                                 }
                             }
+
+                            // Allow services to send requests and notifications now
+                            var safeLanguageServer = (SafeLanguageServer)serviceProvider.GetService<ISafeLanguageServer>();
+                            safeLanguageServer.SetReady();
+
+                            return Task.CompletedTask;
                         });
             }).ConfigureAwait(false);
 

src/PowerShellEditorServices/Server/PsesServiceCollectionExtensions.cs

@@ -17,14 +17,16 @@ public static IServiceCollection AddPsesLanguageServices(
             this IServiceCollection collection,
             HostStartupInfo hostStartupInfo)
         {
-            return collection.AddSingleton<WorkspaceService>()
+            return collection
+                .AddSingleton<ISafeLanguageServer, SafeLanguageServer>()
+                .AddSingleton<WorkspaceService>()
                 .AddSingleton<SymbolsService>()
                 .AddSingleton<ConfigurationService>()
                 .AddSingleton<PowerShellContextService>(
                     (provider) =>
                         PowerShellContextService.Create(
                             provider.GetService<ILoggerFactory>(),
-                            provider.GetService<OmniSharp.Extensions.LanguageServer.Protocol.Server.ILanguageServerFacade>(),
+                            provider.GetService<ISafeLanguageServer>(),
                             hostStartupInfo))
                 .AddSingleton<TemplateService>()
                 .AddSingleton<EditorOperationsService>()
@@ -34,7 +36,7 @@ public static IServiceCollection AddPsesLanguageServices(
                     {
                         var extensionService = new ExtensionService(
                             provider.GetService<PowerShellContextService>(),
-                            provider.GetService<OmniSharp.Extensions.LanguageServer.Protocol.Server.ILanguageServerFacade>());
+                            provider.GetService<ISafeLanguageServer>());
                         extensionService.InitializeAsync(
                             serviceProvider: provider,
                             editorOperations: provider.GetService<EditorOperationsService>())

src/PowerShellEditorServices/Server/SafeLanguageServer.cs

@@ -0,0 +1,233 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+using System;
+using System.Collections.Concurrent;
+using System.Threading;
+using System.Threading.Tasks;
+using MediatR;
+using Newtonsoft.Json.Linq;
+using OmniSharp.Extensions.JsonRpc;
+using OmniSharp.Extensions.LanguageServer.Protocol.Server;
+
+namespace Microsoft.PowerShell.EditorServices.Server
+{
+    /// <summary>
+    /// An LSP server that ensures that client/server initialization has occurred
+    /// before sending or receiving any other notifications or requests.
+    /// </summary>
+    internal interface ISafeLanguageServer : IResponseRouter
+    {
+        ITextDocumentLanguageServer TextDocument { get; }
+
+        IClientLanguageServer Client { get; }
+
+        IGeneralLanguageServer General { get; }
+
+        IWindowLanguageServer Window { get; }
+
+        IWorkspaceLanguageServer Workspace { get; }
+    }
+
+    /// <summary>
+    /// An implementation around Omnisharp's LSP server to ensure
+    /// messages are not sent before initialization has completed.
+    /// </summary>
+    internal class SafeLanguageServer : ISafeLanguageServer
+    {
+        private readonly ILanguageServerFacade _languageServer;
+
+        private readonly AsyncLatch _serverReady;
+
+        private readonly ConcurrentQueue<Action> _notificationQueue;
+
+        public SafeLanguageServer(ILanguageServerFacade languageServer)
+        {
+            _languageServer = languageServer;
+            _serverReady = new AsyncLatch();
+            _notificationQueue = new ConcurrentQueue<Action>();
+        }
+
+        public ITextDocumentLanguageServer TextDocument
+        {
+            get
+            {
+                _serverReady.Wait();
+                return _languageServer.TextDocument;
+            }
+        }
+
+        public IClientLanguageServer Client
+        {
+            get
+            {
+                _serverReady.Wait();
+                return _languageServer.Client;
+            }
+        }
+
+        public IGeneralLanguageServer General
+        {
+            get
+            {
+                _serverReady.Wait();
+                return _languageServer.General;
+            }
+        }
+
+        public IWindowLanguageServer Window
+        {
+            get
+            {
+                _serverReady.Wait();
+                return _languageServer.Window;
+            }
+        }
+
+        public IWorkspaceLanguageServer Workspace
+        {
+            get
+            {
+                _serverReady.Wait();
+                return _languageServer.Workspace;
+            }
+        }
+
+        public void SetReady()
+        {
+            _serverReady.Open();
+
+            // Send any pending notifications now
+            while (_notificationQueue.TryDequeue(out Action notifcationAction))
+            {
+                notifcationAction();
+            }
+        }
+
+        public void SendNotification(string method)
+        {
+            if (!_serverReady.IsReady)
+            {
+                _notificationQueue.Enqueue(() => _languageServer.SendNotification(method));
+                return;
+            }
+
+            _languageServer.SendNotification(method);
+        }
+
+        public void SendNotification<T>(string method, T @params)
+        {
+            if (!_serverReady.IsReady)
+            {
+                _notificationQueue.Enqueue(() => _languageServer.SendNotification(method, @params));
+                return;
+            }
+
+            _languageServer.SendNotification(method, @params);
+        }
+
+        public void SendNotification(IRequest request)
+        {
+            if (!_serverReady.IsReady)
+            {
+                _notificationQueue.Enqueue(() => _languageServer.SendNotification(request));
+                return;
+            }
+
+            _languageServer.SendNotification(request);
+        }
+
+        public IResponseRouterReturns SendRequest<T>(string method, T @params)
+        {
+            _serverReady.Wait();
+            return _languageServer.SendRequest(method, @params);
+        }
+
+        public IResponseRouterReturns SendRequest(string method)
+        {
+            _serverReady.Wait();
+            return _languageServer.SendRequest(method);
+        }
+
+        public async Task<TResponse> SendRequest<TResponse>(IRequest<TResponse> request, CancellationToken cancellationToken)
+        {
+            await _serverReady.WaitAsync();
+            return await _languageServer.SendRequest(request, cancellationToken);
+        }
+
+        public bool TryGetRequest(long id, out string method, out TaskCompletionSource<JToken> pendingTask)
+        {
+            if (!_serverReady.IsReady)
+            {
+                method = default;
+                pendingTask = default;
+                return false;
+            }
+
+            return _languageServer.TryGetRequest(id, out method, out pendingTask);
+        }
+
+        /// <summary>
+        /// Implements a latch (a monotonic manual reset event that starts in the blocking state)
+        /// that can be waited on synchronously or asynchronously without wasting thread resources.
+        /// </summary>
+        private class AsyncLatch
+        {
+            private readonly ManualResetEvent _resetEvent;
+
+            private readonly Task _awaitLatchOpened;
+
+            private volatile bool _isOpen;
+
+            public AsyncLatch()
+            {
+                _resetEvent = new ManualResetEvent(/* start in blocking state */ initialState: false);
+                _awaitLatchOpened = CreateLatchOpenedAwaiterTask(_resetEvent);
+                _isOpen = false;
+            }
+
+            public bool IsReady => _isOpen;
+
+            public void Wait() => _resetEvent.WaitOne();
+
+            public Task WaitAsync() => _awaitLatchOpened;
+
+            public void Open()
+            {
+                // Unblocks the reset event
+                _resetEvent.Set();
+                _isOpen = true;
+            }
+
+            private static Task CreateLatchOpenedAwaiterTask(WaitHandle handle)
+            {
+                var tcs = new TaskCompletionSource<object>();
+
+                // In a dedicated waiter thread, wait for the reset event and then set the task completion source
+                // to turn the reset event wait into a task.
+                // From https://stackoverflow.com/a/18766131.
+                RegisteredWaitHandle registration = ThreadPool.RegisterWaitForSingleObject(handle, (state, timedOut) =>
+                {
+                    ((TaskCompletionSource<object>)state).TrySetResult(result: null);
+                }, tcs, Timeout.Infinite, executeOnlyOnce: true);
+
+                // Register an action to unregister the registration when the reset event task has completed.
+                EnsureWaitHandleUnregistered(tcs.Task, registration);
+
+                return tcs.Task;
+            }
+
+            private static async Task EnsureWaitHandleUnregistered(Task task, RegisteredWaitHandle handle)
+            {
+                try
+                {
+                    await task;
+                }
+                finally
+                {
+                    handle.Unregister(waitObject: null);
+                }
+            }
+        }
+    }
+}