Merge pull request #59 from OmniSharp/fix/52

Fixes #52

Author: david-driscoll

sample/SampleServer/Program.cs

@@ -26,7 +26,7 @@ static async Task MainAsync(string[] args)
             server.AddHandler(new TextDocumentHandler(server));
 
             await server.Initialize();
-            await server.WasShutDown;
+            await server.WaitForExit;
         }
     }
 }

src/Server/Abstractions/IAwaitableTermination.cs

@@ -1,7 +1,10 @@
-namespace OmniSharp.Extensions.LanguageServer.Server.Abstractions
+using System.Threading.Tasks;
+
+namespace OmniSharp.Extensions.LanguageServer.Server.Abstractions
 {
     public interface IAwaitableTermination
     {
-        System.Threading.Tasks.Task WasShutDown { get; }
+        Task WasShutDown { get; }
+        Task WaitForExit { get; }
     }
-}
+}

src/Server/Handlers/ExitHandler.cs

@@ -1,5 +1,6 @@
 using System.Threading.Tasks;
 using OmniSharp.Extensions.LanguageServer.Protocol;
+using OmniSharp.Extensions.LanguageServer.Server.Abstractions;
 
 namespace OmniSharp.Extensions.LanguageServer.Server.Handlers
 {
@@ -12,9 +13,15 @@ public ExitHandler(ShutdownHandler shutdownHandler)
             _shutdownHandler = shutdownHandler;
         }
 
+        private readonly TaskCompletionSource<int> _exitedSource = new TaskCompletionSource<int>(TaskContinuationOptions.LongRunning);
+        public Task WaitForExit => _exitedSource.Task;
+
+
         public Task Handle()
         {
-            Exit?.Invoke(_shutdownHandler.ShutdownRequested ? 0 : 1);
+            var result = _shutdownHandler.ShutdownRequested ? 0 : 1;
+            Exit?.Invoke(result);
+            _exitedSource.SetResult(result);
             return Task.CompletedTask;
         }
 

src/Server/Handlers/ShutdownHandler.cs

@@ -5,14 +5,14 @@
 
 namespace OmniSharp.Extensions.LanguageServer.Server.Handlers
 {
-    public class ShutdownHandler : IShutdownHandler, IAwaitableTermination
+    public class ShutdownHandler : IShutdownHandler
     {
         public event ShutdownEventHandler Shutdown;
 
         public bool ShutdownRequested { get; private set; }
 
         private readonly TaskCompletionSource<bool> _shutdownSource = new TaskCompletionSource<bool>(TaskContinuationOptions.LongRunning);
-        Task IAwaitableTermination.WasShutDown => _shutdownSource.Task;
+        public Task WasShutDown => _shutdownSource.Task;
         public Task Handle(object request, CancellationToken token)
         {
             ShutdownRequested = true;

src/Server/LanguageServer.cs

@@ -329,7 +329,8 @@ public TaskCompletionSource<JToken> GetRequest(long id)
             return _responseRouter.GetRequest(id);
         }
 
-        public Task WasShutDown => ((IAwaitableTermination)_shutdownHandler).WasShutDown;
+        public Task WasShutDown => _shutdownHandler.WasShutDown;
+        public Task WaitForExit => _exitHandler.WaitForExit;
 
         public void Dispose()
         {