WIP: New error handling

Common/Advanced/EarlyAccess/IConnectionSendViewer.cs

@@ -3,8 +3,7 @@
 using System.Threading.Tasks;
 
 namespace MirrorSharp.Advanced.EarlyAccess {
-    internal interface IConnectionSendViewer
-    {
+    internal interface IConnectionSendViewer {
         Task ViewDuringSendAsync(string messageTypeName, ReadOnlyMemory<byte> message, IWorkSession session, CancellationToken cancellationToken);
     }
-}
+}

Common/Internal/Connection.cs

@@ -11,40 +11,39 @@
 using MirrorSharp.Internal.Results;
 
 namespace MirrorSharp.Internal {
-    internal class Connection : ICommandResultSender, IDisposable {
+    internal class Connection : IConnection, ICommandResultSender {
         public static int InputBufferSize => 4096;
 
-        private readonly ArrayPool<byte> _bufferPool;
+        private readonly ArrayPool<byte> _inputBufferPool;
         private readonly IConnectionSendViewer? _sendViewer;
         private readonly WebSocket _socket;
         private readonly WorkSession _session;
         private readonly ImmutableArray<ICommandHandler> _handlers;
         private readonly byte[] _inputBuffer;
 
-        private readonly FastUtf8JsonWriter _messageWriter;
+        private readonly ConnectionMessageWriter _messageWriter;
         private readonly IConnectionOptions? _options;
         private readonly IExceptionLogger? _exceptionLogger;
 
-        private string? _currentMessageTypeName;
-
         public Connection(
             WebSocket socket,
             WorkSession session,
+            ArrayPool<byte> inputBufferPool,
             ImmutableArray<ICommandHandler> handlers,
-            ArrayPool<byte> bufferPool,
+            ConnectionMessageWriter messageWriter,
             IConnectionSendViewer? sendViewer,
             IExceptionLogger? exceptionLogger,
             IConnectionOptions? options
         ) {
             _socket = socket;
             _session = session;
             _handlers = handlers;
-            _messageWriter = new FastUtf8JsonWriter(bufferPool);
+            _messageWriter = messageWriter;
             _options = options;
             _sendViewer = sendViewer;
             _exceptionLogger = exceptionLogger;
-            _bufferPool = bufferPool;
-            _inputBuffer = bufferPool.Rent(InputBufferSize);
+            _inputBufferPool = inputBufferPool;
+            _inputBuffer = inputBufferPool.Rent(InputBufferSize);
         }
 
         public bool IsConnected => _socket.State == WebSocketState.Open;
@@ -138,23 +137,19 @@ private ICommandHandler ResolveHandler(byte commandId) {
         }
 
         private Task SendErrorAsync(string message, CancellationToken cancellationToken) {
-            var writer = StartJsonMessage("error");
-            writer.WriteProperty("message", message);
+            _messageWriter.WriteErrorStart(message);
             return SendJsonMessageAsync(cancellationToken);
         }
 
-        private FastUtf8JsonWriter StartJsonMessage(string messageTypeName) {
-            _messageWriter.Reset();
-            _messageWriter.WriteStartObject();
-            _messageWriter.WriteProperty("type", messageTypeName);
-            _currentMessageTypeName = messageTypeName;
-            return _messageWriter;
-        }
-
         private Task SendJsonMessageAsync(CancellationToken cancellationToken) {
-            _messageWriter.WriteEndObject();
+            _messageWriter.WriteMessageEnd();
 
-            var viewTask = _sendViewer?.ViewDuringSendAsync(_currentMessageTypeName!, _messageWriter.WrittenSegment, _session, cancellationToken);
+            var viewTask = _sendViewer?.ViewDuringSendAsync(
+                _messageWriter.CurrentMessageTypeName!,
+                _messageWriter.WrittenSegment,
+                _session,
+                cancellationToken
+            );
             var sendTask = _socket.SendAsync(
                 _messageWriter.WrittenSegment,
                 WebSocketMessageType.Text, true, cancellationToken
@@ -172,12 +167,24 @@ private async Task WhenAll(Task first, Task second) {
         }
 
         public void Dispose() {
-            _bufferPool.Return(_inputBuffer);
-            _messageWriter.Dispose();
-            _session.Dispose();
+            try {
+                try {
+                    _inputBufferPool.Return(_inputBuffer);
+                }
+                finally {
+                    _messageWriter.Dispose();
+                }
+            }
+            finally {
+                _session.Dispose();
+            }
+        }
+
+        IFastJsonWriter ICommandResultSender.StartJsonMessage(string messageTypeName) {
+            _messageWriter.WriteMessageStart(messageTypeName);
+            return _messageWriter.JsonWriter;
         }
 
-        IFastJsonWriter ICommandResultSender.StartJsonMessage(string messageTypeName) => StartJsonMessage(messageTypeName);
         Task ICommandResultSender.SendJsonMessageAsync(CancellationToken cancellationToken) => SendJsonMessageAsync(cancellationToken);
     }
 }

Common/Internal/ConnectionMessageWriter.cs

@@ -0,0 +1,37 @@
+using System;
+using MirrorSharp.Advanced;
+
+namespace MirrorSharp.Internal {
+    internal class ConnectionMessageWriter : IDisposable {
+        private readonly FastUtf8JsonWriter _jsonWriter;
+        private string? _currentMessageTypeName;
+
+        public ConnectionMessageWriter(FastUtf8JsonWriter jsonWriter) {
+            _jsonWriter = jsonWriter;
+        }
+
+        public void WriteMessageStart(string messageTypeName) {
+            _jsonWriter.Reset();
+            _jsonWriter.WriteStartObject();
+            _jsonWriter.WriteProperty("type", messageTypeName);
+            _currentMessageTypeName = messageTypeName;
+        }
+
+        public void WriteErrorStart(string message) {
+            WriteMessageStart("error");
+            _jsonWriter.WriteProperty("message", message);
+        }
+
+        public void WriteMessageEnd() {
+            _jsonWriter.WriteEndObject();
+        }
+
+        public ArraySegment<byte> WrittenSegment => _jsonWriter.WrittenSegment;
+        public FastUtf8JsonWriter JsonWriter => _jsonWriter;
+        public string? CurrentMessageTypeName => _currentMessageTypeName;
+
+        public void Dispose() {
+            _jsonWriter.Dispose();
+        }
+    }
+}

Common/Internal/FastUtf8JsonWriter.cs

@@ -36,7 +36,7 @@ public FastUtf8JsonWriter(ArrayPool<byte> bufferPool) {
             _encoder = Encoding.UTF8.GetEncoder();
         }
 
-        public ArraySegment<byte> WrittenSegment => new ArraySegment<byte>(_buffer, 0, _position);
+        public ArraySegment<byte> WrittenSegment => new (_buffer, 0, _position);
 
         public void WriteStartObject() {
             WriteStartValue();

Common/Internal/IConnection.cs

@@ -0,0 +1,11 @@
+using System;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace MirrorSharp.Internal {
+    internal interface IConnection : IDisposable {
+        bool IsConnected { get; }
+
+        Task ReceiveAndProcessAsync(CancellationToken cancellationToken);
+    }
+}

Common/Internal/IWorkSessionTracker.cs

@@ -0,0 +1,5 @@
+namespace MirrorSharp.Internal {
+    internal interface IWorkSessionTracker {
+        void TrackNewWorkSession(WorkSession session);
+    }
+}

Common/Internal/ImmutableExtensionServices.cs

@@ -9,14 +9,16 @@ public ImmutableExtensionServices(
             IRoslynSourceTextGuard? roslynSourceTextGuard,
             IRoslynCompilationGuard? roslynCompilationGuard,
             IConnectionSendViewer? connectionSendViewer,
-            IExceptionLogger? exceptionLogger
+            IExceptionLogger? exceptionLogger,
+            IWorkSessionTracker? sessionTracker
         ) {
             SetOptionsFromClient = setOptionsFromClient;
             SlowUpdate = slowUpdate;
             RoslynSourceTextGuard = roslynSourceTextGuard;
             RoslynCompilationGuard = roslynCompilationGuard;
             ConnectionSendViewer = connectionSendViewer;
             ExceptionLogger = exceptionLogger;
+            SessionTracker = sessionTracker;
         }
 
         public ISetOptionsFromClientExtension? SetOptionsFromClient { get; }
@@ -25,5 +27,6 @@ public ImmutableExtensionServices(
         public IRoslynCompilationGuard? RoslynCompilationGuard { get; }
         public IConnectionSendViewer? ConnectionSendViewer { get; }
         public IExceptionLogger? ExceptionLogger { get; }
+        public IWorkSessionTracker? SessionTracker { get; }
     }
 }

Common/Internal/MiddlewareBase.cs

@@ -92,10 +92,19 @@ Task ExecuteAsync(char commandId, string command) {
 
         protected async Task WebSocketLoopAsync(WebSocket socket, CancellationToken cancellationToken) {
             WorkSession? session = null;
-            Connection? connection = null;
+            FastUtf8JsonWriter? messageJsonWriter = null;
+            ConnectionMessageWriter? messageWriter = null;
+            IConnection? connection = null;
             try {
-                session = StartWorkSession();
-                connection = new Connection(socket, session, _handlers, ArrayPool<byte>.Shared, _extensions.ConnectionSendViewer, _extensions.ExceptionLogger, _options);
+                messageJsonWriter = new FastUtf8JsonWriter(ArrayPool<byte>.Shared);
+                messageWriter = new ConnectionMessageWriter(messageJsonWriter);
+                try {
+                    session = StartWorkSession();
+                    connection = new Connection(socket, session, ArrayPool<byte>.Shared, _handlers, messageWriter, _extensions.ConnectionSendViewer, _extensions.ExceptionLogger, _options);
+                }
+                catch (Exception ex) {
+                    connection = new StartupFailedConnection(socket, ex, ArrayPool<byte>.Shared, messageWriter, _options);
+                }
 
                 while (connection.IsConnected) {
                     try {
@@ -111,7 +120,17 @@ protected async Task WebSocketLoopAsync(WebSocket socket, CancellationToken canc
                     connection.Dispose();
                 }
                 else {
-                    session?.Dispose();
+                    try {
+                        if (messageWriter != null) {
+                            messageWriter.Dispose();
+                        }
+                        else {
+                            messageJsonWriter?.Dispose();
+                        }
+                    }
+                    finally {
+                        session?.Dispose();
+                    }
                 }
             }
         }

Common/Internal/Roslyn/RoslynLanguageBase.cs

@@ -46,16 +46,27 @@ IRoslynLanguageOptions options
         }
 
         private CompositionHost CreateCompositionHost(string featuresAssemblyName, string workspacesAssemblyName) {
+            var roslynInternalsAssembly = RoslynInternals.GetInternalsAssemblySlow();
             var types = new[] {
                 RoslynAssemblies.MicrosoftCodeAnalysisWorkspaces,
                 RoslynAssemblies.MicrosoftCodeAnalysisFeatures,
                 Assembly.Load(new AssemblyName(featuresAssemblyName)),
                 Assembly.Load(new AssemblyName(workspacesAssemblyName)),
-                RoslynInternals.GetInternalsAssemblySlow()
+                roslynInternalsAssembly
             }.SelectMany(a => a.DefinedTypes).Where(ShouldConsiderForHostServices);
 
             var configuration = new ContainerConfiguration().WithParts(types);
-            return configuration.CreateContainer();
+            try {
+                return configuration.CreateContainer();
+            }
+            catch (MissingMethodException ex) {
+                throw new Exception(
+                    $"Failed to initialize MirrorSharp {GetType().Name}.{Environment.NewLine}" +
+                    $"  Microsoft.CodeAnalysis {RoslynAssemblies.MicrosoftCodeAnalysis.GetName().Version}.{Environment.NewLine}" +
+                    $"  {roslynInternalsAssembly.GetName().Name}",
+                    ex
+                );
+            }
         }
 
         protected virtual bool ShouldConsiderForHostServices(Type type)

Common/Internal/StartupFailedConnection.cs

@@ -0,0 +1,69 @@
+
+using System;
+using System.Buffers;
+using System.Net.WebSockets;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace MirrorSharp.Internal;
+
+internal class StartupFailedConnection : IConnection {
+    public static int InputBufferSize => 4096;
+
+    private readonly ArrayPool<byte> _bufferPool;
+    private readonly WebSocket _socket;
+    private readonly Exception _startupException;
+    private readonly byte[] _inputBuffer;
+
+    private readonly ConnectionMessageWriter _messageWriter;
+    private readonly IConnectionOptions? _options;
+
+    public StartupFailedConnection(
+        WebSocket socket,
+        Exception startupException,
+        ArrayPool<byte> bufferPool,
+        ConnectionMessageWriter messageWriter,
+        IConnectionOptions? options
+    ) {
+        _socket = socket;
+        _startupException = startupException;
+        _messageWriter = messageWriter;
+        _options = options;
+        _bufferPool = bufferPool;
+        _inputBuffer = bufferPool.Rent(InputBufferSize);
+    }
+
+    public bool IsConnected => _socket.State == WebSocketState.Open;
+
+    public async Task ReceiveAndProcessAsync(CancellationToken cancellationToken) {
+        var first = await _socket.ReceiveAsync(new ArraySegment<byte>(_inputBuffer), cancellationToken).ConfigureAwait(false);
+        if (first.MessageType == WebSocketMessageType.Close) {
+            await _socket.CloseAsync(first.CloseStatus ?? WebSocketCloseStatus.Empty, first.CloseStatusDescription, cancellationToken).ConfigureAwait(false);
+            return;
+        }
+
+        if (!first.EndOfMessage)
+            await ReceiveToEndAsync(cancellationToken).ConfigureAwait(false);
+
+        var error = (_options?.IncludeExceptionDetails ?? false)
+            ? _startupException.ToString()
+            : "A server error has occurred during startup.";
+
+        _messageWriter.WriteErrorStart(error);
+        _messageWriter.WriteMessageEnd();
+        await _socket.SendAsync(
+            _messageWriter.WrittenSegment,
+            WebSocketMessageType.Text, true, cancellationToken
+        );
+    }
+
+    private async Task ReceiveToEndAsync(CancellationToken cancellationToken) {
+        while (!(await _socket.ReceiveAsync(new ArraySegment<byte>(_inputBuffer), cancellationToken).ConfigureAwait(false)).EndOfMessage) {
+        }
+    }
+
+    public void Dispose() {
+        _bufferPool.Return(_inputBuffer);
+        _messageWriter.Dispose();
+    }
+}

Internal.Roslyn411/MirrorSharpDiagnosticAnalyzerService.cs

@@ -21,7 +21,7 @@ internal class MirrorSharpDiagnosticAnalyzerService : IDiagnosticAnalyzerService
         Task<ImmutableArray<DiagnosticData>> IDiagnosticAnalyzerService.GetDiagnosticsForIdsAsync(Solution solution, ProjectId? projectId, DocumentId? documentId, ImmutableHashSet<string>? diagnosticIds, Func<DiagnosticAnalyzer, bool>? shouldIncludeAnalyzer, Func<Project, DocumentId?, IReadOnlyList<DocumentId>>? getDocumentIds, bool includeSuppressedDiagnostics, bool includeLocalDocumentDiagnostics, bool includeNonLocalDocumentDiagnostics, CancellationToken cancellationToken) => throw new NotSupportedException();
         Task<ImmutableArray<DiagnosticData>> IDiagnosticAnalyzerService.GetDiagnosticsForSpanAsync(TextDocument document, TextSpan? range, Func<string, bool>? shouldIncludeDiagnostic, bool includeCompilerDiagnostics, bool includeSuppressedDiagnostics, ICodeActionRequestPriorityProvider priorityProvider, DiagnosticKind diagnosticKind, bool isExplicit, CancellationToken cancellationToken) => throw new NotSupportedException();
         Task<ImmutableArray<DiagnosticData>> IDiagnosticAnalyzerService.GetProjectDiagnosticsForIdsAsync(Solution solution, ProjectId? projectId, ImmutableHashSet<string>? diagnosticIds, Func<DiagnosticAnalyzer, bool>? shouldIncludeAnalyzer, bool includeSuppressedDiagnostics, bool includeNonLocalDocumentDiagnostics, CancellationToken cancellationToken) => throw new NotSupportedException();
-        void IDiagnosticAnalyzerService.RequestDiagnosticRefresh() => throw new NotImplementedException();
+        void IDiagnosticAnalyzerService.RequestDiagnosticRefresh() => throw new NotSupportedException();
         Task<(ImmutableArray<DiagnosticData> diagnostics, bool upToDate)> IDiagnosticAnalyzerService.TryGetDiagnosticsForSpanAsync(TextDocument document, TextSpan range, Func<string, bool>? shouldIncludeDiagnostic, bool includeSuppressedDiagnostics, ICodeActionRequestPriorityProvider priorityProvider, DiagnosticKind diagnosticKind, bool isExplicit, CancellationToken cancellationToken) => throw new NotSupportedException();
     }
 }

Testing/Internal/HandlerTestArgument.cs

@@ -39,6 +39,10 @@ public static implicit operator HandlerTestArgument(byte[][] data) {
             return new HandlerTestArgument(data);
         }
 
+        public byte[] ToBytes(char commandId) {
+
+        }
+
         public AsyncData ToAsyncData(char commandId) {
             var nextIndex = 1;
 

Testing/Internal/TestMiddleware.cs

@@ -0,0 +1,16 @@
+using System.Net.WebSockets;
+using System.Threading.Tasks;
+using System.Threading;
+using MirrorSharp.Internal;
+
+namespace MirrorSharp.Testing.Internal {
+    internal class TestMiddleware : MiddlewareBase {
+        public TestMiddleware(LanguageManager languageManager, IMiddlewareOptions options, ImmutableExtensionServices extensions)
+            : base(languageManager, options, extensions) {
+        }
+
+        public Task WebSocketLoopAsync(TestWebSocket socket, CancellationToken cancellationToken) {
+            return base.WebSocketLoopAsync(socket, cancellationToken);
+        }
+    }
+}