Implement a notification queue so sends don't block

Author: rjmholt

src/PowerShellEditorServices/Server/SafeLanguageServer.cs

@@ -1,6 +1,8 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
+using System;
+using System.Collections.Concurrent;
 using System.Threading;
 using System.Threading.Tasks;
 using MediatR;
@@ -10,6 +12,10 @@
 
 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; }
@@ -23,12 +29,25 @@ internal interface ISafeLanguageServer : IResponseRouter
  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
@@ -77,29 +96,44 @@ public IWorkspaceLanguageServer Workspace
  public void SetReady()
  {
  _serverReady.Open();
- }
 
- public SafeLanguageServer(ILanguageServerFacade languageServer)
- {
- _languageServer = languageServer;
- _serverReady = new AsyncLatch();
+ // Send any pending notifications now
+ while (_notificationQueue.TryDequeue(out Action notifcationAction))
+ {
+ notifcationAction();
+ }
  }
 
  public void SendNotification(string method)
  {
- _serverReady.Wait();
+ if (!_serverReady.IsReady)
+ {
+ _notificationQueue.Enqueue(() => _languageServer.SendNotification(method));
+ return;
+ }
+
  _languageServer.SendNotification(method);
  }
 
  public void SendNotification<T>(string method, T @params)
  {
- _serverReady.Wait();
+ if (!_serverReady.IsReady)
+ {
+ _notificationQueue.Enqueue(() => _languageServer.SendNotification(method, @params));
+ return;
+ }
+
  _languageServer.SendNotification(method, @params);
  }
 
  public void SendNotification(IRequest request)
  {
- _serverReady.Wait();
+ if (!_serverReady.IsReady)
+ {
+ _notificationQueue.Enqueue(() => _languageServer.SendNotification(request));
+ return;
+ }
+
  _languageServer.SendNotification(request);
  }
 
@@ -123,7 +157,7 @@ public async Task<TResponse> SendRequest<TResponse>(IRequest<TResponse> request,
 
  public bool TryGetRequest(long id, out string method, out TaskCompletionSource<JToken> pendingTask)
  {
- if (!_serverReady.TryWait())
+ if (!_serverReady.IsReady)
  {
  method = default;
  pendingTask = default;
@@ -133,6 +167,10 @@ public bool TryGetRequest(long id, out string method, out TaskCompletionSource<J
  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;
@@ -148,12 +186,12 @@ public AsyncLatch()
  _isOpen = false;
  }
 
+ public bool IsReady => _isOpen;
+
  public void Wait() => _resetEvent.WaitOne();
 
  public Task WaitAsync() => _awaitLatchOpened;
 
- public bool TryWait() => _isOpen;
-
  public void Open()
  {
  // Unblocks the reset event