Fix the flaky keep alive tests

closes #4301

src/Compilers/CSharp/Test/CommandLine/TouchedFileLoggingTests.cs

@@ -215,7 +215,7 @@ public void TrivialMetadataCaching()
                 filelist.Add(source1);
                 var outWriter = new StringWriter();
                 var cmd = new CSharpCompilerServer(
-                    new DesktopCompilerServerHost(Guid.NewGuid().ToString()),
+                    new DesktopCompilerServerHost(),
                     new[] { "/nologo", "/touchedfiles:" + touchedBase, source1 },
                     null,
                     _baseDirectory,

src/Compilers/Server/PortableServer/CoreClrCompilerServerHost.cs

@@ -17,17 +17,15 @@ internal sealed class CoreClrCompilerServerHost : CompilerServerHost
     {
         private readonly Func<string, MetadataReferenceProperties, PortableExecutableReference> _assemblyReferenceProvider = (path, properties) => new CachingMetadataReference(path, properties);
         private readonly IAnalyzerAssemblyLoader _analyzerAssemblyLoader = CoreClrAnalyzerAssemblyLoader.CreateAndSetDefault();
-        private readonly TcpListener _listener;
 
         public override IAnalyzerAssemblyLoader AnalyzerAssemblyLoader => _analyzerAssemblyLoader;
 
         public override Func<string, MetadataReferenceProperties, PortableExecutableReference> AssemblyReferenceProvider => _assemblyReferenceProvider;
 
-        internal CoreClrCompilerServerHost(IPEndPoint endPoint, string clientDirectory)
+        internal CoreClrCompilerServerHost(string clientDirectory)
             :base(clientDirectory : clientDirectory, sdkDirectory: null)
         {
-            _listener = new TcpListener(endPoint);
-            _listener.Start();
+
         }
 
         public override bool CheckAnalyzers(string baseDirectory, ImmutableArray<CommandLineAnalyzerReference> analyzers)
@@ -36,51 +34,9 @@ public override bool CheckAnalyzers(string baseDirectory, ImmutableArray<Command
             return analyzers.Length == 0;
         }
 
-        public override async Task<IClientConnection> CreateListenTask(CancellationToken cancellationToken)
-        {
-            var tcpClient = await _listener.AcceptTcpClientAsync().ConfigureAwait(true);
-            return new TcpClientConnection(tcpClient);
-        }
-
         public override void Log(string message)
         {
             // BTODO: Do we need this anymore? 
         }
-
-        private sealed class TcpClientConnection : IClientConnection
-        {
-            private readonly string _identifier = Guid.NewGuid().ToString();
-            private readonly TcpClient _client;
-
-            public string LoggingIdentifier => _identifier;
-
-            internal TcpClientConnection(TcpClient client)
-            {
-                _client = client;
-            }
-
-            public Task<BuildRequest> ReadBuildRequest(CancellationToken cancellationToken)
-            {
-                return BuildRequest.ReadAsync(_client.GetStream(), cancellationToken);
-            }
-
-            public Task WriteBuildResponse(BuildResponse response, CancellationToken cancellationToken)
-            {
-                return response.WriteAsync(_client.GetStream(), cancellationToken);
-            }
-
-            public async Task CreateMonitorDisconnectTask(CancellationToken cancellationToken)
-            {
-                while (_client.Connected && !cancellationToken.IsCancellationRequested)
-                {
-                    await Task.Delay(TimeSpan.FromSeconds(1)).ConfigureAwait(false);
-                }
-            }
-
-            public void Close()
-            {
-                _client.Dispose();
-            }
-        }
     }
 }

src/Compilers/Server/PortableServer/PortableServer.csproj

@@ -65,6 +65,7 @@
     </Compile>
     <Compile Include="CoreClrCompilerServerHost.cs" />
     <Compile Include="Program.cs" />
+    <Compile Include="TcpClientConnection.cs" />
   </ItemGroup>
   <ItemGroup>
     <None Include="project.json" />

src/Compilers/Server/PortableServer/Program.cs

@@ -13,9 +13,9 @@ public static void Main(string[] args)
             var ipAddress = IPAddress.Parse("127.0.0.1");
             var endPoint = new IPEndPoint(ipAddress, port: 12000);
             var clientDirectory = AppContext.BaseDirectory;
-            var compilerHost = new CoreClrCompilerServerHost(endPoint, clientDirectory);
-            var compilerRequestHandler = new CompilerRequestHandler(compilerHost);
-            var serverDispatcher = new ServerDispatcher(compilerHost, compilerRequestHandler, new EmptyDiagnosticListener());
+            var compilerHost = new CoreClrCompilerServerHost(clientDirectory);
+            var connectionHost = new TcpClientConnectionHost(compilerHost, endPoint);
+            var serverDispatcher = new ServerDispatcher(connectionHost);
             serverDispatcher.ListenAndDispatchConnections(keepAlive: null, cancellationToken: CancellationToken.None);
         }
     }

src/Compilers/Server/PortableServer/TcpClientConnection.cs

+// Copyright (c) Microsoft.  All Rights Reserved.  Licensed under the Apache License, Version 2.0.  See License.txt in the project root for license information.
+
+using Microsoft.CodeAnalysis.CommandLine;
+using System;
+using System.Collections.Generic;
+using System.Collections.Immutable;
+using System.Linq;
+using System.Net;
+using System.Net.Sockets;
+using System.Text;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Microsoft.CodeAnalysis.CompilerServer
+{
+    internal sealed class TcpClientConnectionHost : IClientConnectionHost
+    {
+        private readonly ICompilerServerHost _compilerServerHost;
+        private readonly TcpListener _listener;
+        private int _connectionCount;
+
+        internal TcpClientConnectionHost(ICompilerServerHost compilerServerHost, IPEndPoint endPoint)
+        {
+            _compilerServerHost = compilerServerHost;
+            _listener = new TcpListener(endPoint);
+            _listener.Start();
+        }
+
+        public async Task<IClientConnection> CreateListenTask(CancellationToken cancellationToken)
+        {
+            var tcpClient = await _listener.AcceptTcpClientAsync().ConfigureAwait(true);
+            return new TcpClientConnection(_compilerServerHost, tcpClient, _connectionCount++.ToString());
+        }
+
+        private sealed class TcpClientConnection : ClientConnection
+        {
+            private readonly TcpClient _client;
+
+            internal TcpClientConnection(ICompilerServerHost compilerServerHost, TcpClient client, string loggingIdentifier) : base(compilerServerHost, loggingIdentifier, client.GetStream())
+            {
+                _client = client;
+            }
+
+            public override void Close()
+            {
+                _client.Dispose();
+            }
+
+            protected override Task CreateMonitorDisconnectTask(CancellationToken cancellationToken)
+            {
+                return Task.Delay(-1, cancellationToken);
+            }
+        }
+    }
+}
+

src/Compilers/Server/ServerShared/CompilerRequestHandler.cs

@@ -29,56 +29,6 @@ public RunRequest(string language, string currentDirectory, string libDirectory,
         }
     }
 
-    internal interface IRequestHandler
-    {
-        BuildResponse HandleRequest(BuildRequest req, CancellationToken cancellationToken);
-    }
-
-    internal sealed class CompilerRequestHandler : IRequestHandler
-    {
-        private readonly ICompilerServerHost _compilerServerHost;
-
-        internal CompilerRequestHandler(ICompilerServerHost compilerServerHost)
-        {
-            _compilerServerHost = compilerServerHost;
-        }
-
-        /// <summary>
-        /// An incoming request as occurred. This is called on a new thread to handle
-        /// the request.
-        /// </summary>
-        public BuildResponse HandleRequest(BuildRequest buildRequest, CancellationToken cancellationToken)
-        {
-            var request = BuildProtocolUtil.GetRunRequest(buildRequest);
-            CommonCompiler compiler;
-            if (!_compilerServerHost.TryCreateCompiler(request, out compiler))
-            {
-                // We can't do anything with a request we don't know about. 
-                _compilerServerHost.Log($"Got request with id '{request.Language}'");
-                return new CompletedBuildResponse(-1, false, "", "");
-            }
-
-            _compilerServerHost.Log($"CurrentDirectory = '{request.CurrentDirectory}'");
-            _compilerServerHost.Log($"LIB = '{request.LibDirectory}'");
-            for (int i = 0; i < request.Arguments.Length; ++i)
-            {
-                _compilerServerHost.Log($"Argument[{i}] = '{request.Arguments[i]}'");
-            }
-
-            bool utf8output = compiler.Arguments.Utf8Output;
-            if (!_compilerServerHost.CheckAnalyzers(request.CurrentDirectory, compiler.Arguments.AnalyzerReferences))
-            {
-                return new AnalyzerInconsistencyBuildResponse();
-            }
-
-            _compilerServerHost.Log($"****Running {request.Language} compiler...");
-            TextWriter output = new StringWriter(CultureInfo.InvariantCulture);
-            int returnCode = compiler.Run(output, cancellationToken);
-            _compilerServerHost.Log($"****{request.Language} Compilation complete.\r\n****Return code: {returnCode}\r\n****Output:\r\n{output.ToString()}\r\n");
-            return new CompletedBuildResponse(returnCode, utf8output, output.ToString(), "");
-        }
-    }
-
     internal abstract class CompilerServerHost : ICompilerServerHost
     {
         public abstract IAnalyzerAssemblyLoader AnalyzerAssemblyLoader { get; }
@@ -101,8 +51,6 @@ protected CompilerServerHost(string clientDirectory, string sdkDirectory)
             SdkDirectory = sdkDirectory;
         }
 
-        public abstract Task<IClientConnection> CreateListenTask(CancellationToken cancellationToken);
-
         public abstract bool CheckAnalyzers(string baseDirectory, ImmutableArray<CommandLineAnalyzerReference> analyzers);
 
         public abstract void Log(string message);

src/Compilers/Server/ServerShared/Connection.cs

@@ -2,11 +2,8 @@
 
 using Roslyn.Utilities;
 using System;
-using System.Diagnostics;
+using System.Globalization;
 using System.IO;
-using System.IO.Pipes;
-using System.Runtime.CompilerServices;
-using System.Security.Principal;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.CommandLine;
@@ -43,6 +40,12 @@ internal enum CompletionReason
         /// the results could be provided to them.  
         /// </summary>
         ClientDisconnect,
+
+        /// <summary>
+        /// There was an unhandled exception processing the result.
+        /// BTODO: need to handle this is the server.
+        /// </summary>
+        ClientException,
     }
 
     /// <summary>
@@ -50,20 +53,32 @@ internal enum CompletionReason
     /// from when the client connects to it, until the request is finished or abandoned.
     /// A new task is created to actually service the connection and do the operation.
     /// </summary>
-    internal sealed class Connection
+    internal abstract class ClientConnection : IClientConnection
     {
-        private readonly IClientConnection _clientConnection;
-        private readonly IRequestHandler _handler;
+        private readonly ICompilerServerHost _compilerServerHost;
         private readonly string _loggingIdentifier;
+        private readonly Stream _stream;
 
-        public Connection(IClientConnection clientConnection, IRequestHandler handler)
+        public string LoggingIdentifier => _loggingIdentifier;
+
+        public ClientConnection(ICompilerServerHost compilerServerHost, string loggingIdentifier, Stream stream)
         {
-            _clientConnection = clientConnection;
-            _loggingIdentifier = clientConnection.LoggingIdentifier;
-            _handler = handler;
+            _compilerServerHost = compilerServerHost;
+            _loggingIdentifier = loggingIdentifier;
+            _stream = stream;
         }
 
-        public async Task<ConnectionData> ServeConnection(CancellationToken cancellationToken = default(CancellationToken))
+        /// <summary>
+        /// Returns a Task that resolves if the client stream gets disconnected.  
+        /// </summary>
+        protected abstract Task CreateMonitorDisconnectTask(CancellationToken cancellationToken);
+
+        /// <summary>
+        /// Close the connection.  Can be called multiple times.
+        /// </summary>
+        public abstract void Close();
+
+        public async Task<ConnectionData> HandleConnection(CancellationToken cancellationToken)
         {
             try
             {
@@ -71,7 +86,7 @@ public async Task<ConnectionData> ServeConnection(CancellationToken cancellation
                 try
                 {
                     Log("Begin reading request.");
-                    request = await _clientConnection.ReadBuildRequest(cancellationToken).ConfigureAwait(false);
+                    request = await BuildRequest.ReadAsync(_stream, cancellationToken).ConfigureAwait(false);
                     Log("End reading request.");
                 }
                 catch (Exception e)
@@ -85,7 +100,7 @@ public async Task<ConnectionData> ServeConnection(CancellationToken cancellation
                 // Kick off both the compilation and a task to monitor the pipe for closing.  
                 var buildCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
                 var compilationTask = ServeBuildRequest(request, buildCts.Token);
-                var monitorTask = _clientConnection.CreateMonitorDisconnectTask(buildCts.Token);
+                var monitorTask = CreateMonitorDisconnectTask(buildCts.Token);
                 await Task.WhenAny(compilationTask, monitorTask).ConfigureAwait(false);
 
                 // Do an 'await' on the completed task, preference being compilation, to force
@@ -98,7 +113,7 @@ public async Task<ConnectionData> ServeConnection(CancellationToken cancellation
                     try
                     {
                         Log("Begin writing response.");
-                        await _clientConnection.WriteBuildResponse(response, cancellationToken).ConfigureAwait(false);
+                        await response.WriteAsync(_stream, cancellationToken).ConfigureAwait(false);
                         reason = CompletionReason.Completed;
                         Log("End writing response.");
                     }
@@ -119,7 +134,7 @@ public async Task<ConnectionData> ServeConnection(CancellationToken cancellation
             }
             finally
             {
-                _clientConnection.Close();
+                Close();
             }
         }
 
@@ -149,7 +164,7 @@ public async Task<ConnectionData> ServeConnection(CancellationToken cancellation
             return timeout;
         }
 
-        private Task<BuildResponse> ServeBuildRequest(BuildRequest request, CancellationToken cancellationToken)
+        protected virtual Task<BuildResponse> ServeBuildRequest(BuildRequest request, CancellationToken cancellationToken)
         {
             return Task.Run(() =>
             {
@@ -157,7 +172,7 @@ private Task<BuildResponse> ServeBuildRequest(BuildRequest request, Cancellation
                 {
                     // Do the compilation
                     Log("Begin compilation");
-                    BuildResponse response = _handler.HandleRequest(request, cancellationToken);
+                    BuildResponse response = ServeBuildRequestCore(request, cancellationToken);
                     Log("End compilation");
                     return response;
                 }
@@ -168,6 +183,37 @@ private Task<BuildResponse> ServeBuildRequest(BuildRequest request, Cancellation
             });
         }
 
+        private BuildResponse ServeBuildRequestCore(BuildRequest buildRequest, CancellationToken cancellationToken)
+        {
+            var request = BuildProtocolUtil.GetRunRequest(buildRequest);
+            CommonCompiler compiler;
+            if (!_compilerServerHost.TryCreateCompiler(request, out compiler))
+            {
+                // We can't do anything with a request we don't know about. 
+                _compilerServerHost.Log($"Got request with id '{request.Language}'");
+                return new CompletedBuildResponse(-1, false, "", "");
+            }
+
+            _compilerServerHost.Log($"CurrentDirectory = '{request.CurrentDirectory}'");
+            _compilerServerHost.Log($"LIB = '{request.LibDirectory}'");
+            for (int i = 0; i < request.Arguments.Length; ++i)
+            {
+                _compilerServerHost.Log($"Argument[{i}] = '{request.Arguments[i]}'");
+            }
+
+            bool utf8output = compiler.Arguments.Utf8Output;
+            if (!_compilerServerHost.CheckAnalyzers(request.CurrentDirectory, compiler.Arguments.AnalyzerReferences))
+            {
+                return new AnalyzerInconsistencyBuildResponse();
+            }
+
+            _compilerServerHost.Log($"****Running {request.Language} compiler...");
+            TextWriter output = new StringWriter(CultureInfo.InvariantCulture);
+            int returnCode = compiler.Run(output, cancellationToken);
+            _compilerServerHost.Log($"****{request.Language} Compilation complete.\r\n****Return code: {returnCode}\r\n****Output:\r\n{output.ToString()}\r\n");
+            return new CompletedBuildResponse(returnCode, utf8output, output.ToString(), "");
+        }
+
         private void Log(string message)
         {
             CompilerServerLogger.Log("Client {0}: {1}", _loggingIdentifier, message);
@@ -177,5 +223,6 @@ private void LogException(Exception e, string message)
         {
             CompilerServerLogger.LogException(e, string.Format("Client {0}: {1}", _loggingIdentifier, message));
         }
+
     }
 }

src/Compilers/Server/ServerShared/IClientConnection.cs

@@ -20,26 +20,13 @@ internal interface IClientConnection
         /// A value which can be used to identify this connection for logging purposes only.  It has 
         /// no guarantee of uniqueness.  
         /// </summary>
-        string LoggingIdentifier
-        {
-            get;
-        }
-
-        /// <summary>
-        /// Read the <see cref="BuildRequest"/> object from the client connection.
-        /// </summary>
-        Task<BuildRequest> ReadBuildRequest(CancellationToken cancellationToken);
-
-        /// <summary>
-        /// Write the <see cref="BuildRequest"/> object to the client connection.
-        /// </summary>
-        Task WriteBuildResponse(BuildResponse response, CancellationToken cancellationToken);
+        string LoggingIdentifier { get; }
 
         /// <summary>
-        /// Create a <see cref="Task"/> object which will complete if the client connection is broken
-        /// by the client.
+        /// Server the connection and return the result.
         /// </summary>
-        Task CreateMonitorDisconnectTask(CancellationToken cancellationToken);
+        /// <param name="cancellationToken"></param>
+        Task<ConnectionData> HandleConnection(CancellationToken cancellationToken);
 
         /// <summary>
         /// Close the underlying client connection.

src/Compilers/Server/ServerShared/ICompilerServerHost.cs

@@ -12,13 +12,11 @@ namespace Microsoft.CodeAnalysis.CompilerServer
 {
     internal interface ICompilerServerHost
     {
+        // BTODO: how many of these are needed anymore? 
         IAnalyzerAssemblyLoader AnalyzerAssemblyLoader { get; }
         Func<string, MetadataReferenceProperties, PortableExecutableReference> AssemblyReferenceProvider { get; }
         bool TryCreateCompiler(RunRequest request, out CommonCompiler compiler);
         bool CheckAnalyzers(string baseDirectory, ImmutableArray<CommandLineAnalyzerReference> analyzers);
         void Log(string message);
-
-        // BTODO: Move this to a new interface
-        Task<IClientConnection> CreateListenTask(CancellationToken cancellationToken);
     }
 }

src/Compilers/Server/ServerShared/ServerDispatcher.cs

@@ -13,15 +13,15 @@
 
 namespace Microsoft.CodeAnalysis.CompilerServer
 {
+    internal interface IClientConnectionHost
+    {
+        Task<IClientConnection> CreateListenTask(CancellationToken cancellationToken);
+    }
+
     /// <summary>
-    /// This class handles the named pipe creation, listening, thread creation,
-    /// and so forth. When a request comes in, it is dispatched on a new thread
-    /// to the <see cref="IRequestHandler"/> interface. The request handler does the actual
-    /// compilation. This class itself has no dependencies on the compiler.
+    /// This class manages the connections, timeout and general scheduling of the client 
+    /// requests.  
     /// </summary>
-    /// <remarks>
-    /// One instance of this is created per process.
-    /// </remarks>
     internal sealed class ServerDispatcher
     {
         /// <summary>
@@ -35,15 +35,13 @@ internal sealed class ServerDispatcher
         /// </summary>
         internal static readonly TimeSpan GCTimeout = TimeSpan.FromSeconds(30);
 
-        private readonly ICompilerServerHost _compilerServerHost;
-        private readonly IRequestHandler _handler;
+        private readonly IClientConnectionHost _clientConnectionHost;
         private readonly IDiagnosticListener _diagnosticListener;
 
-        internal ServerDispatcher(ICompilerServerHost compilerServerHost, IRequestHandler handler, IDiagnosticListener diagnosticListener)
+        internal ServerDispatcher(IClientConnectionHost clientConnectionHost, IDiagnosticListener diagnosticListener = null)
         {
-            _compilerServerHost = compilerServerHost;
-            _handler = handler;
-            _diagnosticListener = diagnosticListener;
+            _clientConnectionHost = clientConnectionHost;
+            _diagnosticListener = diagnosticListener ?? new EmptyDiagnosticListener();
         }
 
         /// <summary>
@@ -71,7 +69,7 @@ public void ListenAndDispatchConnections(TimeSpan? keepAlive, CancellationToken
                     Debug.Assert(listenCancellationTokenSource == null);
                     Debug.Assert(timeoutTask == null);
                     listenCancellationTokenSource = new CancellationTokenSource();
-                    listenTask = _compilerServerHost.CreateListenTask(listenCancellationTokenSource.Token);
+                    listenTask = _clientConnectionHost.CreateListenTask(listenCancellationTokenSource.Token);
                 }
 
                 // If there are no active clients running then the server needs to be in a timeout mode.
@@ -86,7 +84,7 @@ public void ListenAndDispatchConnections(TimeSpan? keepAlive, CancellationToken
                 // If there is a connection event that has highest priority. 
                 if (listenTask.IsCompleted && !cancellationToken.IsCancellationRequested)
                 {
-                    var connectionTask = CreateHandleConnectionTask(listenTask, _handler, cancellationToken);
+                    var connectionTask = HandleClientConnection(listenTask, cancellationToken);
                     connectionList.Add(connectionTask);
                     listenTask = null;
                     listenCancellationTokenSource = null;
@@ -213,13 +211,12 @@ private void ChangeKeepAlive(TimeSpan? value, ref TimeSpan? keepAlive, ref bool
         /// will never fail.  It will always produce a <see cref="ConnectionData"/> value.  Connection errors
         /// will end up being represented as <see cref="CompletionReason.ClientDisconnect"/>
         /// </summary>
-        internal static async Task<ConnectionData> CreateHandleConnectionTask(Task<IClientConnection> clientConnectionTask, IRequestHandler handler, CancellationToken cancellationToken)
+        internal static async Task<ConnectionData> HandleClientConnection(Task<IClientConnection> clientConnectionTask, CancellationToken cancellationToken = default(CancellationToken))
         {
-            Connection connection;
+            IClientConnection clientConnection;
             try
             {
-                var clientConnection = await clientConnectionTask.ConfigureAwait(false);
-                connection = new Connection(clientConnection, handler);
+                clientConnection = await clientConnectionTask.ConfigureAwait(false);
             }
             catch (Exception ex)
             {
@@ -229,7 +226,15 @@ internal static async Task<ConnectionData> CreateHandleConnectionTask(Task<IClie
                 return new ConnectionData(CompletionReason.CompilationNotStarted);
             }
 
-            return await connection.ServeConnection(cancellationToken).ConfigureAwait(false);
+            try
+            {
+                return await clientConnection.HandleConnection(cancellationToken).ConfigureAwait(false);
+            }
+            catch (Exception ex)
+            {
+                CompilerServerLogger.LogException(ex, "Error handling connection");
+                return new ConnectionData(CompletionReason.ClientException);
+            }
         }
     }
 }

src/Compilers/Server/VBCSCompiler/DesktopCompilerServerHost.cs

@@ -18,30 +18,25 @@ namespace Microsoft.CodeAnalysis.CompilerServer
 {
     internal sealed class DesktopCompilerServerHost : CompilerServerHost
     {
-        // Size of the buffers to use
-        private const int PipeBufferSize = 0x10000;  // 64K
-
         private static readonly IAnalyzerAssemblyLoader s_analyzerLoader = new ShadowCopyAnalyzerAssemblyLoader(Path.Combine(Path.GetTempPath(), "VBCSCompiler", "AnalyzerAssemblyLoader"));
 
         // Caches are used by C# and VB compilers, and shared here.
         private static readonly Func<string, MetadataReferenceProperties, PortableExecutableReference> s_assemblyReferenceProvider = (path, properties) => new CachingMetadataReference(path, properties);
 
-        private readonly string _pipeName;
-
         public override IAnalyzerAssemblyLoader AnalyzerAssemblyLoader => s_analyzerLoader;
 
         public override Func<string, MetadataReferenceProperties, PortableExecutableReference> AssemblyReferenceProvider => s_assemblyReferenceProvider;
 
-        internal DesktopCompilerServerHost(string pipeName)
-            : this(pipeName, AppDomain.CurrentDomain.BaseDirectory, RuntimeEnvironment.GetRuntimeDirectory())
+        internal DesktopCompilerServerHost()
+            : this(AppDomain.CurrentDomain.BaseDirectory, RuntimeEnvironment.GetRuntimeDirectory())
         {
 
         }
 
-        internal DesktopCompilerServerHost(string pipeName, string clientDirectory, string sdkDirectory)
+        internal DesktopCompilerServerHost(string clientDirectory, string sdkDirectory)
             : base(clientDirectory, sdkDirectory)
         {
-            _pipeName = pipeName;
+
         }
 
         public override bool CheckAnalyzers(string baseDirectory, ImmutableArray<CommandLineAnalyzerReference> analyzers)
@@ -54,149 +49,5 @@ public override void Log(string message)
             CompilerServerLogger.Log(message);
         }
 
-        public override async Task<IClientConnection> CreateListenTask(CancellationToken cancellationToken)
-        {
-            var pipeStream = await CreateListenTaskCore(cancellationToken).ConfigureAwait(false);
-            return new NamedPipeClientConnection(pipeStream);
-        }
-
-        /// <summary>
-        /// Creates a Task that waits for a client connection to occur and returns the connected 
-        /// <see cref="NamedPipeServerStream"/> object.  Throws on any connection error.
-        /// </summary>
-        /// <param name="cancellationToken">Used to cancel the connection sequence.</param>
-        private async Task<NamedPipeServerStream> CreateListenTaskCore(CancellationToken cancellationToken)
-        {
-            // Create the pipe and begin waiting for a connection. This 
-            // doesn't block, but could fail in certain circumstances, such
-            // as Windows refusing to create the pipe for some reason 
-            // (out of handles?), or the pipe was disconnected before we 
-            // starting listening.
-            NamedPipeServerStream pipeStream = ConstructPipe(_pipeName);
-
-            // Unfortunately the version of .Net we are using doesn't support the WaitForConnectionAsync
-            // method.  When it is available it should absolutely be used here.  In the meantime we
-            // have to deal with the idea that this WaitForConnection call will block a thread
-            // for a significant period of time.  It is unadvisable to do this to a thread pool thread 
-            // hence we will use an explicit thread here.
-            var listenSource = new TaskCompletionSource<NamedPipeServerStream>();
-            var listenTask = listenSource.Task;
-            var listenThread = new Thread(() =>
-            {
-                try
-                {
-                    CompilerServerLogger.Log("Waiting for new connection");
-                    pipeStream.WaitForConnection();
-                    CompilerServerLogger.Log("Pipe connection detected.");
-
-                    if (Environment.Is64BitProcess || MemoryHelper.IsMemoryAvailable())
-                    {
-                        CompilerServerLogger.Log("Memory available - accepting connection");
-                        listenSource.SetResult(pipeStream);
-                        return;
-                    }
-
-                    try
-                    {
-                        pipeStream.Close();
-                    }
-                    catch
-                    {
-                        // Okay for Close failure here.  
-                    }
-
-                    listenSource.SetException(new Exception("Insufficient resources to process new connection."));
-                }
-                catch (Exception ex)
-                {
-                    listenSource.SetException(ex);
-                }
-            });
-            listenThread.Start();
-
-            // Create a tasks that waits indefinitely (-1) and completes only when cancelled.
-            var waitCancellationTokenSource = new CancellationTokenSource();
-            var waitTask = Task.Delay(
-                Timeout.Infinite,
-                CancellationTokenSource.CreateLinkedTokenSource(waitCancellationTokenSource.Token, cancellationToken).Token);
-            await Task.WhenAny(listenTask, waitTask).ConfigureAwait(false);
-            if (listenTask.IsCompleted)
-            {
-                waitCancellationTokenSource.Cancel();
-                return await listenTask.ConfigureAwait(false);
-            }
-
-            // The listen operation was cancelled.  Close the pipe stream throw a cancellation exception to
-            // simulate the cancel operation.
-            waitCancellationTokenSource.Cancel();
-            try
-            {
-                pipeStream.Close();
-            }
-            catch
-            {
-                // Okay for Close failure here.
-            }
-
-            throw new OperationCanceledException();
-        }
-
-        /// <summary>
-        /// Creates a Task representing the processing of the new connection.  This will return a task that
-        /// will never fail.  It will always produce a <see cref="ConnectionData"/> value.  Connection errors
-        /// will end up being represented as <see cref="CompletionReason.ClientDisconnect"/>
-        /// </summary>
-        internal static async Task<ConnectionData> CreateHandleConnectionTask(Task<NamedPipeServerStream> pipeStreamTask, IRequestHandler handler, CancellationToken cancellationToken)
-        {
-            Connection connection;
-            try
-            {
-                var pipeStream = await pipeStreamTask.ConfigureAwait(false);
-                var clientConnection = new NamedPipeClientConnection(pipeStream);
-                connection = new Connection(clientConnection, handler);
-            }
-            catch (Exception ex)
-            {
-                // Unable to establish a connection with the client.  The client is responsible for
-                // handling this case.  Nothing else for us to do here.
-                CompilerServerLogger.LogException(ex, "Error creating client named pipe");
-                return new ConnectionData(CompletionReason.CompilationNotStarted);
-            }
-
-            return await connection.ServeConnection(cancellationToken).ConfigureAwait(false);
-        }
-
-        /// <summary>
-        /// Create an instance of the pipe. This might be the first instance, or a subsequent instance.
-        /// There always needs to be an instance of the pipe created to listen for a new client connection.
-        /// </summary>
-        /// <returns>The pipe instance or throws an exception.</returns>
-        private NamedPipeServerStream ConstructPipe(string pipeName)
-        {
-            CompilerServerLogger.Log("Constructing pipe '{0}'.", pipeName);
-
-            SecurityIdentifier identifier = WindowsIdentity.GetCurrent().Owner;
-            PipeSecurity security = new PipeSecurity();
-
-            // Restrict access to just this account.  
-            PipeAccessRule rule = new PipeAccessRule(identifier, PipeAccessRights.ReadWrite | PipeAccessRights.CreateNewInstance, AccessControlType.Allow);
-            security.AddAccessRule(rule);
-            security.SetOwner(identifier);
-
-            NamedPipeServerStream pipeStream = new NamedPipeServerStream(
-                pipeName,
-                PipeDirection.InOut,
-                NamedPipeServerStream.MaxAllowedServerInstances, // Maximum connections.
-                PipeTransmissionMode.Byte,
-                PipeOptions.Asynchronous | PipeOptions.WriteThrough,
-                PipeBufferSize, // Default input buffer
-                PipeBufferSize, // Default output buffer
-                security,
-                HandleInheritability.None);
-
-            CompilerServerLogger.Log("Successfully constructed pipe '{0}'.", pipeName);
-
-            return pipeStream;
-        }
     }
 }

src/Compilers/Server/VBCSCompiler/NamedPipeClientConnection.cs

@@ -10,83 +10,164 @@
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.CommandLine;
+using System.Security.AccessControl;
 
 namespace Microsoft.CodeAnalysis.CompilerServer
 {
-    internal sealed class NamedPipeClientConnection : IClientConnection
+    internal sealed class NamedPipeClientConnectionHost : IClientConnectionHost
     {
-        private readonly NamedPipeServerStream _pipeStream;
+        // Size of the buffers to use: 64K
+        private const int PipeBufferSize = 0x10000;
 
-        // This is a value used for logging only, do not depend on this value
-        private readonly string _loggingIdentifier;
-        private static int s_lastLoggingIdentifier;
+        private readonly ICompilerServerHost _compilerServerHost;
+        private readonly string _pipeName;
+        private int _loggingIdentifier;
 
-        internal NamedPipeClientConnection(NamedPipeServerStream pipeStream)
+        internal NamedPipeClientConnectionHost(ICompilerServerHost compilerServerHost, string pipeName)
         {
-            _pipeStream = pipeStream;
-            _loggingIdentifier = Interlocked.Increment(ref s_lastLoggingIdentifier).ToString();
+            _compilerServerHost = compilerServerHost;
+            _pipeName = pipeName;
         }
 
-        public string LoggingIdentifier
+        public async Task<IClientConnection> CreateListenTask(CancellationToken cancellationToken)
         {
-            get { return _loggingIdentifier; }
+            var pipeStream = await CreateListenTaskCore(cancellationToken).ConfigureAwait(false);
+            return new NamedPipeClientConnection(_compilerServerHost, _loggingIdentifier++.ToString(), pipeStream);
         }
 
         /// <summary>
-        /// The IsConnected property on named pipes does not detect when the client has disconnected
-        /// if we don't attempt any new I/O after the client disconnects. We start an async I/O here
-        /// which serves to check the pipe for disconnection. 
-        ///
-        /// This will return true if the pipe was disconnected.
+        /// Creates a Task that waits for a client connection to occur and returns the connected 
+        /// <see cref="NamedPipeServerStream"/> object.  Throws on any connection error.
         /// </summary>
-        private async Task<bool> CreateMonitorDisconnectTaskCore(CancellationToken cancellationToken)
+        /// <param name="cancellationToken">Used to cancel the connection sequence.</param>
+        private async Task<NamedPipeServerStream> CreateListenTaskCore(CancellationToken cancellationToken)
+        {
+            // Create the pipe and begin waiting for a connection. This 
+            // doesn't block, but could fail in certain circumstances, such
+            // as Windows refusing to create the pipe for some reason 
+            // (out of handles?), or the pipe was disconnected before we 
+            // starting listening.
+            NamedPipeServerStream pipeStream = ConstructPipe(_pipeName);
+
+            // Unfortunately the version of .Net we are using doesn't support the WaitForConnectionAsync
+            // method.  When it is available it should absolutely be used here.  In the meantime we
+            // have to deal with the idea that this WaitForConnection call will block a thread
+            // for a significant period of time.  It is unadvisable to do this to a thread pool thread 
+            // hence we will use an explicit thread here.
+            var listenSource = new TaskCompletionSource<NamedPipeServerStream>();
+            var listenTask = listenSource.Task;
+            var listenThread = new Thread(() =>
             {
-            var buffer = SpecializedCollections.EmptyBytes;
+                try
+                {
+                    CompilerServerLogger.Log("Waiting for new connection");
+                    pipeStream.WaitForConnection();
+                    CompilerServerLogger.Log("Pipe connection detected.");
 
-            while (!cancellationToken.IsCancellationRequested && _pipeStream.IsConnected)
+                    if (!ClientAndOurIdentitiesMatch(pipeStream))
                     {
-                // Wait a second before trying again
-                await Task.Delay(1000, cancellationToken).ConfigureAwait(false);
+                        var exception = new Exception("Client identity does not match server identity.");
+                        listenSource.SetException(exception);
+                        return;
+                    }
+
+                    if (Environment.Is64BitProcess || MemoryHelper.IsMemoryAvailable())
+                    {
+                        CompilerServerLogger.Log("Memory available - accepting connection");
+                        listenSource.SetResult(pipeStream);
+                        return;
+                    }
 
                     try
                     {
-                    CompilerServerLogger.Log("Pipe {0}: Before poking pipe.", _loggingIdentifier);
-                    await _pipeStream.ReadAsync(buffer, 0, 0, cancellationToken).ConfigureAwait(false);
-                    CompilerServerLogger.Log("Pipe {0}: After poking pipe.", _loggingIdentifier);
+                        pipeStream.Close();
                     }
-                catch (Exception e)
+                    catch
                     {
-                    // It is okay for this call to fail.  Errors will be reflected in the 
-                    // IsConnected property which will be read on the next iteration of the 
-                    // loop
-                    var msg = string.Format("Pipe {0}: Error poking pipe.", _loggingIdentifier);
-                    CompilerServerLogger.LogException(e, msg);
+                        // Okay for Close failure here.  
                     }
+
+                    listenSource.SetException(new Exception("Insufficient resources to process new connection."));
+                }
+                catch (Exception ex)
+                {
+                    listenSource.SetException(ex);
+                }
+            });
+            listenThread.Start();
+
+            // Create a tasks that waits indefinitely (-1) and completes only when cancelled.
+            var waitCancellationTokenSource = new CancellationTokenSource();
+            var waitTask = Task.Delay(
+                Timeout.Infinite,
+                CancellationTokenSource.CreateLinkedTokenSource(waitCancellationTokenSource.Token, cancellationToken).Token);
+            await Task.WhenAny(listenTask, waitTask).ConfigureAwait(false);
+            if (listenTask.IsCompleted)
+            {
+                waitCancellationTokenSource.Cancel();
+                return await listenTask.ConfigureAwait(false);
+            }
+
+            // The listen operation was cancelled.  Close the pipe stream throw a cancellation exception to
+            // simulate the cancel operation.
+            waitCancellationTokenSource.Cancel();
+            try
+            {
+                pipeStream.Close();
+            }
+            catch
+            {
+                // Okay for Close failure here.
+            }
+
+            throw new OperationCanceledException();
         }
 
-            return !_pipeStream.IsConnected;
+        /// <summary>
+        /// Create an instance of the pipe. This might be the first instance, or a subsequent instance.
+        /// There always needs to be an instance of the pipe created to listen for a new client connection.
+        /// </summary>
+        /// <returns>The pipe instance or throws an exception.</returns>
+        private NamedPipeServerStream ConstructPipe(string pipeName)
+        {
+            CompilerServerLogger.Log("Constructing pipe '{0}'.", pipeName);
+
+            SecurityIdentifier identifier = WindowsIdentity.GetCurrent().Owner;
+            PipeSecurity security = new PipeSecurity();
+
+            // Restrict access to just this account.  
+            PipeAccessRule rule = new PipeAccessRule(identifier, PipeAccessRights.ReadWrite | PipeAccessRights.CreateNewInstance, AccessControlType.Allow);
+            security.AddAccessRule(rule);
+            security.SetOwner(identifier);
+
+            NamedPipeServerStream pipeStream = new NamedPipeServerStream(
+                pipeName,
+                PipeDirection.InOut,
+                NamedPipeServerStream.MaxAllowedServerInstances, // Maximum connections.
+                PipeTransmissionMode.Byte,
+                PipeOptions.Asynchronous | PipeOptions.WriteThrough,
+                PipeBufferSize, // Default input buffer
+                PipeBufferSize, // Default output buffer
+                security,
+                HandleInheritability.None);
+
+            CompilerServerLogger.Log("Successfully constructed pipe '{0}'.", pipeName);
+
+            return pipeStream;
         }
 
         /// <summary>
         /// Does the client of "pipeStream" have the same identity and elevation as we do?
         /// </summary>
-        private bool ClientAndOurIdentitiesMatch()
+        private static bool ClientAndOurIdentitiesMatch(NamedPipeServerStream pipeStream)
         {
             var serverIdentity = GetIdentity(impersonating: false);
 
             Tuple<string, bool> clientIdentity = null;
-            _pipeStream.RunAsClient(() => { clientIdentity = GetIdentity(impersonating: true); });
-
-            CompilerServerLogger.Log(
-                "Pipe {0}: Server identity = '{1}', server elevation='{2}'.",
-                _loggingIdentifier,
-                serverIdentity.Item1,
-                serverIdentity.Item2.ToString());
-            CompilerServerLogger.Log(
-                "Pipe {0}: Client identity = '{1}', client elevation='{2}'.",
-                _loggingIdentifier,
-                clientIdentity.Item1,
-                clientIdentity.Item2.ToString());
+            pipeStream.RunAsClient(() => { clientIdentity = GetIdentity(impersonating: true); });
+
+            CompilerServerLogger.Log($"Server identity = '{serverIdentity.Item1}', server elevation='{serverIdentity.Item2}'.");
+            CompilerServerLogger.Log($"Client identity = '{clientIdentity.Item1}', client elevation='{serverIdentity.Item2}'.");
 
             return
                 StringComparer.OrdinalIgnoreCase.Equals(serverIdentity.Item1, clientIdentity.Item1) &&
@@ -103,43 +184,66 @@ private bool ClientAndOurIdentitiesMatch()
             var elevatedToAdmin = currentPrincipal.IsInRole(WindowsBuiltInRole.Administrator);
             return Tuple.Create(currentIdentity.Name, elevatedToAdmin);
         }
+    }
+
+    internal sealed class NamedPipeClientConnection : ClientConnection
+    {
+        private readonly NamedPipeServerStream _pipeStream;
+
+        internal NamedPipeClientConnection(ICompilerServerHost compilerServerHost, string loggingIdentifier, NamedPipeServerStream pipeStream)
+            : base(compilerServerHost, loggingIdentifier, pipeStream)
+        {
+            _pipeStream = pipeStream;
+        }
+
+        /// <summary>
+        /// The IsConnected property on named pipes does not detect when the client has disconnected
+        /// if we don't attempt any new I/O after the client disconnects. We start an async I/O here
+        /// which serves to check the pipe for disconnection. 
+        ///
+        /// This will return true if the pipe was disconnected.
+        /// </summary>
+        protected override async Task CreateMonitorDisconnectTask(CancellationToken cancellationToken)
+        {
+            var buffer = SpecializedCollections.EmptyBytes;
 
-        public void Close()
+            while (!cancellationToken.IsCancellationRequested && _pipeStream.IsConnected)
             {
-            CompilerServerLogger.Log("Pipe {0}: Closing.", _loggingIdentifier);
+                // Wait a second before trying again
+                await Task.Delay(1000, cancellationToken).ConfigureAwait(false);
+
                 try
                 {
-                _pipeStream.Close();
+                    CompilerServerLogger.Log($"Pipe {LoggingIdentifier}: Before poking pipe.");
+                    await _pipeStream.ReadAsync(buffer, 0, 0, cancellationToken).ConfigureAwait(false);
+                    CompilerServerLogger.Log($"Pipe {LoggingIdentifier}: After poking pipe.");
                 }
                 catch (Exception e)
                 {
-                // The client connection failing to close isn't fatal to the server process.  It is simply a client
-                // for which we can no longer communicate and that's okay because the Close method indicates we are
-                // done with the client already.
-                var msg = string.Format("Pipe {0}: Error closing pipe.", _loggingIdentifier);
+                    // It is okay for this call to fail.  Errors will be reflected in the 
+                    // IsConnected property which will be read on the next iteration of the 
+                    // loop
+                    var msg = string.Format($"Pipe {LoggingIdentifier}: Error poking pipe.");
                     CompilerServerLogger.LogException(e, msg);
                 }
             }
-
-        public Task CreateMonitorDisconnectTask(CancellationToken cancellationToken)
-        {
-            return CreateMonitorDisconnectTaskCore(cancellationToken);
         }
 
-        public async Task<BuildRequest> ReadBuildRequest(CancellationToken cancellationToken)
+        public override void Close()
         {
-            var buildRequest = await BuildRequest.ReadAsync(_pipeStream, cancellationToken).ConfigureAwait(false);
-            if (!ClientAndOurIdentitiesMatch())
+            CompilerServerLogger.Log($"Pipe {LoggingIdentifier}: Closing.");
+            try
             {
-                throw new Exception("Client identity does not match server identity.");
-            }
-
-            return buildRequest;
+                _pipeStream.Close();
             }
-
-        public Task WriteBuildResponse(BuildResponse response, CancellationToken cancellationToken)
+            catch (Exception e)
             {
-            return response.WriteAsync(_pipeStream, cancellationToken);
+                // The client connection failing to close isn't fatal to the server process.  It is simply a client
+                // for which we can no longer communicate and that's okay because the Close method indicates we are
+                // done with the client already.
+                var msg = string.Format($"Pipe {LoggingIdentifier}: Error closing pipe.");
+                CompilerServerLogger.LogException(e, msg);
+            }
         }
     }
 }

src/Compilers/Server/VBCSCompiler/VBCSCompiler.cs

@@ -28,7 +28,7 @@ public static int Main(string[] args)
 
             // VBCSCompiler is installed in the same directory as csc.exe and vbc.exe which is also the 
             // location of the response files.
-            var compilerExeDirectory = AppDomain.CurrentDomain.BaseDirectory;
+            var clientDirectory = AppDomain.CurrentDomain.BaseDirectory;
 
             // Pipename should be passed as the first and only argument to the server process
             // and it must have the form "-pipename:name". Otherwise, exit with a non-zero
@@ -59,7 +59,7 @@ public static int Main(string[] args)
 
                 try
                 {
-                    return Run(keepAliveTimeout, compilerExeDirectory, pipeName);
+                    return Run(keepAliveTimeout, clientDirectory, pipeName);
                 }
                 finally
                 {
@@ -68,7 +68,7 @@ public static int Main(string[] args)
             }
         }
 
-        private static int Run(TimeSpan? keepAliveTimeout, string compilerExeDirectory, string pipeName)
+        private static int Run(TimeSpan? keepAliveTimeout, string clientDirectory, string pipeName)
         {
             try
             {
@@ -102,15 +102,11 @@ private static int Run(TimeSpan? keepAliveTimeout, string compilerExeDirectory,
             FatalError.Handler = FailFast.OnFatalException;
 
             var sdkDirectory = RuntimeEnvironment.GetRuntimeDirectory();
-            var compilerServerHost = new DesktopCompilerServerHost(pipeName, compilerExeDirectory, sdkDirectory);
-            var dispatcher = new ServerDispatcher(
-                compilerServerHost,
-                new CompilerRequestHandler(compilerServerHost),
-                new EmptyDiagnosticListener());
-
+            var compilerServerHost = new DesktopCompilerServerHost(clientDirectory, sdkDirectory);
+            var clientConnectionHost = new NamedPipeClientConnectionHost(compilerServerHost, pipeName);
+            var dispatcher = new ServerDispatcher(clientConnectionHost, new EmptyDiagnosticListener());
             dispatcher.ListenAndDispatchConnections(keepAliveTimeout);
             return CommonCompiler.Succeeded;
         }
-
     }
 }

src/Compilers/Server/VBCSCompilerTests/ClientConnectionTests.cs

+// Copyright (c) Microsoft.  All Rights Reserved.  Licensed under the Apache License, Version 2.0.  See License.txt in the project root for license information.
+
+using System;
+using System.Collections.Generic;
+using System.Collections.Immutable;
+using System.IO;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.CodeAnalysis.CommandLine;
+using Moq;
+using Roslyn.Test.Utilities;
+using Xunit;
+using System.Runtime.InteropServices;
+
+namespace Microsoft.CodeAnalysis.CompilerServer.UnitTests
+{
+    public class ClientConnectionTests 
+    {
+        private sealed class TestableClientConnection : ClientConnection
+        {
+            internal Stream Stream;
+            internal Func<CancellationToken, Task> CreateMonitorDisconnectTaskFunc;
+            internal Func<BuildRequest, CancellationToken, Task<BuildResponse>> ServeBuildRequestFunc;
+
+            internal TestableClientConnection(ICompilerServerHost compilerServerHost, Stream stream)
+                :base(compilerServerHost, "identifier", stream)
+            {
+                Stream = stream;
+                CreateMonitorDisconnectTaskFunc = ct => Task.Delay(-1, ct);
+            }
+
+            public override void Close()
+            {
+
+            }
+
+            protected override Task CreateMonitorDisconnectTask(CancellationToken cancellationToken)
+            {
+                return CreateMonitorDisconnectTaskFunc(cancellationToken);
+            }
+
+            protected override Task<BuildResponse> ServeBuildRequest(BuildRequest request, CancellationToken cancellationToken)
+            {
+                if (ServeBuildRequestFunc != null)
+                {
+                    return ServeBuildRequestFunc(request, cancellationToken);
+                }
+
+                return base.ServeBuildRequest(request, cancellationToken);
+            }
+        }
+
+        private static readonly BuildRequest s_emptyCSharpBuildRequest = new BuildRequest(
+            1,
+            RequestLanguage.CSharpCompile,
+            ImmutableArray<BuildRequest.Argument>.Empty);
+
+        private static readonly BuildResponse s_emptyBuildResponse = new CompletedBuildResponse(
+            returnCode: 0,
+            utf8output: false,
+            output: string.Empty,
+            errorOutput: string.Empty);
+
+        private static TestableClientConnection CreateConnection(Stream stream, ICompilerServerHost compilerServerHost = null)
+        {
+            compilerServerHost = compilerServerHost ?? new Mock<ICompilerServerHost>().Object;
+            return new TestableClientConnection(compilerServerHost, stream);
+        }
+
+        [Fact]
+        public async Task ReadFailure()
+        {
+            var stream = new Mock<Stream>(MockBehavior.Strict);
+            var connection = CreateConnection(stream.Object);
+            var result = await connection.HandleConnection(CancellationToken.None).ConfigureAwait(true);
+            Assert.Equal(CompletionReason.CompilationNotStarted, result.CompletionReason);
+        }
+
+        /// <summary>
+        /// A failure to write the results to the client is considered a client disconnection.  Any error
+        /// from when the build starts to when the write completes should be handled this way. 
+        /// </summary>
+        [Fact]
+        public async Task WriteError()
+        {
+            var realStream = new MemoryStream();
+            await s_emptyCSharpBuildRequest.WriteAsync(realStream, CancellationToken.None).ConfigureAwait(true);
+            realStream.Position = 0;
+
+            var stream = new Mock<Stream>(MockBehavior.Strict);
+            stream
+                .Setup(x => x.ReadAsync(It.IsAny<byte[]>(), It.IsAny<int>(), It.IsAny<int>(), It.IsAny<CancellationToken>()))
+                .Returns((byte[] array, int start, int length, CancellationToken ct) => Task.FromResult(realStream.Read(array, start, length)));
+
+            var connection = CreateConnection(stream.Object);
+            connection.ServeBuildRequestFunc = delegate { return Task.FromResult(s_emptyBuildResponse); };
+            var connectionData = await connection.HandleConnection(CancellationToken.None).ConfigureAwait(true);
+            Assert.Equal(CompletionReason.ClientDisconnect, connectionData.CompletionReason);
+            Assert.Null(connectionData.KeepAlive);
+        }
+
+        /// <summary>
+        /// Ensure the Connection correctly handles the case where a client disconnects while in the 
+        /// middle of a build event.
+        /// </summary>
+        [Fact]
+        public async Task ClientDisconnectsDuringBuild()
+        {
+            var memoryStream = new MemoryStream();
+            await s_emptyCSharpBuildRequest.WriteAsync(memoryStream, CancellationToken.None).ConfigureAwait(true);
+            memoryStream.Position = 0;
+
+            // Fake a long running build task here that we can validate later on. 
+            var buildTaskSource = new TaskCompletionSource<BuildResponse>();
+            var buildTaskCancellationToken = default(CancellationToken);
+            var clientConnection = CreateConnection(memoryStream);
+            clientConnection.ServeBuildRequestFunc = (req, ct) =>
+            {
+                buildTaskCancellationToken = ct;
+                return buildTaskSource.Task;
+            };
+
+            var readyTaskSource = new TaskCompletionSource<bool>();
+            var monitorTaskSource = new TaskCompletionSource<bool>();
+            clientConnection.CreateMonitorDisconnectTaskFunc = (ct) =>
+            {
+                readyTaskSource.SetResult(true);
+                return monitorTaskSource.Task;
+            };
+
+            var handleTask = clientConnection.HandleConnection(CancellationToken.None);
+
+            // Wait until the monitor task is actually created and running. 
+            await readyTaskSource.Task.ConfigureAwait(false);
+
+            // Now simulate a disconnect by the client.  
+            monitorTaskSource.SetResult(true);
+
+            var connectionData = await handleTask.ConfigureAwait(true);
+            Assert.Equal(CompletionReason.ClientDisconnect, connectionData.CompletionReason);
+            Assert.Null(connectionData.KeepAlive);
+            Assert.True(buildTaskCancellationToken.IsCancellationRequested);
+        }
+    }
+}

src/Compilers/Server/VBCSCompilerTests/CompilerServerApiTest.cs

@@ -16,40 +16,6 @@ namespace Microsoft.CodeAnalysis.CompilerServer.UnitTests
 {
     public class CompilerServerApiTest : TestBase
     {
-        private sealed class TestableClientConnection : IClientConnection
-        {
-            internal readonly string LoggingIdentifier = string.Empty;
-            internal Task<BuildRequest> ReadBuildRequestTask = TaskFromException<BuildRequest>(new Exception());
-            internal Task WriteBuildResponseTask = TaskFromException(new Exception());
-            internal Task MonitorTask = TaskFromException(new Exception());
-            internal Action CloseAction = delegate { };
-
-            string IClientConnection.LoggingIdentifier
-            {
-                get { return LoggingIdentifier; }
-            }
-
-            Task<BuildRequest> IClientConnection.ReadBuildRequest(CancellationToken cancellationToken)
-            {
-                return ReadBuildRequestTask;
-            }
-
-            Task IClientConnection.WriteBuildResponse(BuildResponse response, CancellationToken cancellationToken)
-            {
-                return WriteBuildResponseTask;
-            }
-
-            Task IClientConnection.CreateMonitorDisconnectTask(CancellationToken cancellationToken)
-            {
-                return MonitorTask;
-            }
-
-            void IClientConnection.Close()
-            {
-                CloseAction();
-            }
-        }
-
         private sealed class TestableDiagnosticListener : IDiagnosticListener
         {
             public int ProcessedCount;
@@ -101,6 +67,32 @@ private static Task<T> TaskFromException<T>(Exception e)
             return source.Task;
         }
 
+        private static IClientConnection CreateClientConnection(CompletionReason completionReason, TimeSpan? keepAlive = null)
+        {
+            var task = Task.FromResult(new ConnectionData(completionReason, keepAlive));
+            return CreateClientConnection(task);
+        }
+
+        private static IClientConnection CreateClientConnection(Task<ConnectionData> task)
+        {
+            var connection = new Mock<IClientConnection>();
+            connection
+                .Setup(x => x.HandleConnection(It.IsAny<CancellationToken>()))
+                .Returns(task);
+            return connection.Object;
+        }
+
+        private static IClientConnectionHost CreateClientConnectionHost(params Task<IClientConnection>[] connections)
+        {
+            var host = new Mock<IClientConnectionHost>();
+            var index = 0;
+            host
+                .Setup(x => x.CreateListenTask(It.IsAny<CancellationToken>()))
+                .Returns((CancellationToken ct) => connections[index++]);
+
+            return host.Object;
+        }
+
         private async Task<BuildRequest> CreateBuildRequest(string sourceText, TimeSpan? keepAlive = null)
         {
             var directory = Temp.CreateDirectory();
@@ -136,168 +128,78 @@ private async Task<BuildResponse> RunCSharpCompile(string pipeName, string sourc
             }
         }
 
-        /// <summary>
-        /// This returns an <see cref="IRequestHandler"/> that always returns <see cref="CompletedBuildResponse"/> without
-        /// doing any work.
-        /// </summary>
-        private static Mock<IRequestHandler> CreateNopRequestHandler()
+        private static Mock<IClientConnectionHost> CreateNopClientConnectionHost()
         {
-            var requestHandler = new Mock<IRequestHandler>();
-            requestHandler
-                .Setup(x => x.HandleRequest(It.IsAny<BuildRequest>(), It.IsAny<CancellationToken>()))
-                .Returns(new CompletedBuildResponse(0, utf8output: false, output: string.Empty, errorOutput: string.Empty));
-            return requestHandler;
-        }
-
-        private static Mock<ICompilerServerHost> CreateNopCompilerServerHost()
-        {
-            var host = new Mock<ICompilerServerHost>();
+            var host = new Mock<IClientConnectionHost>();
             host
                 .Setup(x => x.CreateListenTask(It.IsAny<CancellationToken>()))
                 .Returns(new TaskCompletionSource<IClientConnection>().Task);
             return host;
         }
 
-        [Fact]
-        public void NotifyCallBackOnRequestHandlerException()
-        {
-            var clientConnection = new TestableClientConnection();
-            clientConnection.MonitorTask = Task.Delay(-1);
-            clientConnection.ReadBuildRequestTask = Task.FromResult(s_emptyCSharpBuildRequest);
-
-            var ex = new Exception();
-            var handler = new Mock<IRequestHandler>();
-            handler
-                .Setup(x => x.HandleRequest(It.IsAny<BuildRequest>(), It.IsAny<CancellationToken>()))
-                .Throws(ex);
-
-            var invoked = false;
-            FatalError.OverwriteHandler((providedEx) =>
+        private static Task<T> FromException<T>(Exception ex)
         {
-                Assert.Same(ex, providedEx);
-                invoked = true;
-            });
-            var client = new Connection(clientConnection, handler.Object);
-
-            Assert.Throws(typeof(AggregateException), () => client.ServeConnection().Wait());
-            Assert.True(invoked);
+            var source = new TaskCompletionSource<T>();
+            source.SetException(ex);
+            return source.Task;
         }
 
         [Fact]
-        public void ClientDisconnectCancelBuildAndReturnsFailure()
+        public async Task ClientConnectionThrowsHandlingBuild()
         {
-            var clientConnection = new TestableClientConnection();
-            clientConnection.ReadBuildRequestTask = Task.FromResult(s_emptyCSharpBuildRequest);
-
-            var monitorTaskSource = new TaskCompletionSource<bool>();
-            clientConnection.MonitorTask = monitorTaskSource.Task;
-
-            var handler = new Mock<IRequestHandler>();
-            var handlerTaskSource = new TaskCompletionSource<CancellationToken>();
-            var releaseHandlerSource = new TaskCompletionSource<bool>();
-            handler
-                .Setup(x => x.HandleRequest(It.IsAny<BuildRequest>(), It.IsAny<CancellationToken>()))
-                .Callback<BuildRequest, CancellationToken>((_, t) =>
-                {
-                    handlerTaskSource.SetResult(t);
-                    releaseHandlerSource.Task.Wait();
-                })
-                .Returns(s_emptyBuildResponse);
-
-            var client = new Connection(clientConnection, handler.Object);
-            var serveTask = client.ServeConnection();
-
-            // Once this returns we know the Connection object has kicked off a compilation and 
-            // started monitoring the disconnect task.  Can now initiate a disconnect in a known
-            // state.
-            var cancellationToken = handlerTaskSource.Task.Result;
-            monitorTaskSource.SetResult(true);
-
-            Assert.Equal(CompletionReason.ClientDisconnect, serveTask.Result.CompletionReason);
-            Assert.True(cancellationToken.IsCancellationRequested);
-
-            // Now that the asserts are done unblock the "build" long running task.  Have to do this
-            // last to simulate a build which is still running when the client disconnects.
-            releaseHandlerSource.SetResult(true);
-        }
+            var ex = new Exception();
+            var clientConnection = new Mock<IClientConnection>();
+            clientConnection
+                .Setup(x => x.HandleConnection(It.IsAny<CancellationToken>()))
+                .Returns(FromException<ConnectionData>(ex));
 
-        [Fact]
-        public void ReadError()
-        {
-            var handler = new Mock<IRequestHandler>(MockBehavior.Strict);
-            var ex = new Exception("Simulated read error.");
-            var clientConnection = new TestableClientConnection();
-            var calledClose = false;
-            clientConnection.ReadBuildRequestTask = TaskFromException<BuildRequest>(ex);
-            clientConnection.CloseAction = delegate { calledClose = true; };
-
-            var client = new Connection(clientConnection, handler.Object);
-            Assert.Equal(CompletionReason.CompilationNotStarted, client.ServeConnection().Result.CompletionReason);
-            Assert.True(calledClose);
+            var task = Task.FromResult(clientConnection.Object);
+
+            var connectionData = await ServerDispatcher.HandleClientConnection(task).ConfigureAwait(true);
+            Assert.Equal(CompletionReason.ClientException, connectionData.CompletionReason);
+            Assert.Null(connectionData.KeepAlive);
         }
 
-        /// <summary>
-        /// A failure to write the results to the client is considered a client disconnection.  Any error
-        /// from when the build starts to when the write completes should be handled this way. 
-        /// </summary>
         [Fact]
-        public void WriteError()
+        public async Task ClientConnectionThrowsConnecting()
         {
-            var clientConnection = new TestableClientConnection();
-            clientConnection.MonitorTask = Task.Delay(-1);
-            clientConnection.ReadBuildRequestTask = Task.FromResult(s_emptyCSharpBuildRequest);
-            clientConnection.WriteBuildResponseTask = TaskFromException(new Exception());
-            var handler = new Mock<IRequestHandler>();
-            handler
-                .Setup(x => x.HandleRequest(It.IsAny<BuildRequest>(), It.IsAny<CancellationToken>()))
-                .Returns(s_emptyBuildResponse);
-
-            var client = new Connection(clientConnection, handler.Object);
-            Assert.Equal(CompletionReason.ClientDisconnect, client.ServeConnection().Result.CompletionReason);
+            var ex = new Exception();
+            var task = FromException<IClientConnection>(ex);
+            var connectionData = await ServerDispatcher.HandleClientConnection(task).ConfigureAwait(true);
+            Assert.Equal(CompletionReason.CompilationNotStarted, connectionData.CompletionReason);
+            Assert.Null(connectionData.KeepAlive);
         }
 
         [Fact]
         public void KeepAliveNoConnections()
         {
             var keepAlive = TimeSpan.FromSeconds(3);
-            var requestHandler = new Mock<IRequestHandler>(MockBehavior.Strict);
-            var dispatcher = new ServerDispatcher(CreateNopCompilerServerHost().Object, requestHandler.Object, new EmptyDiagnosticListener());
+            var connectionHost = new Mock<IClientConnectionHost>();
+            connectionHost
+                .Setup(x => x.CreateListenTask(It.IsAny<CancellationToken>()))
+                .Returns(new TaskCompletionSource<IClientConnection>().Task);
+
+            var dispatcher = new ServerDispatcher(connectionHost.Object, new EmptyDiagnosticListener());
             var startTime = DateTime.Now;
             dispatcher.ListenAndDispatchConnections(keepAlive);
 
             Assert.True((DateTime.Now - startTime) > keepAlive);
         }
 
-        [Fact]
-        public async Task FailedConnectionShouldCreateFailedConnectionData()
-        {
-            var tcs = new TaskCompletionSource<IClientConnection>();
-            var handler = new Mock<IRequestHandler>(MockBehavior.Strict);
-            var connectionDataTask = ServerDispatcher.CreateHandleConnectionTask(tcs.Task, handler.Object, CancellationToken.None);
-
-            tcs.SetException(new Exception());
-            var connectionData = await connectionDataTask.ConfigureAwait(false);
-            Assert.Equal(CompletionReason.CompilationNotStarted, connectionData.CompletionReason);
-            Assert.Null(connectionData.KeepAlive);
-        }
-
         /// <summary>
         /// Ensure server respects keep alive and shuts down after processing a single connection.
         /// </summary>
-        [Fact(Skip = "https://github.com/dotnet/roslyn/issues/4301")]
-        public async Task KeepAliveAfterSingleConnection()
+        [Fact]
+        public void KeepAliveAfterSingleConnection()
         {
-            var keepAlive = TimeSpan.FromSeconds(1);
+            var connection = CreateClientConnection(CompletionReason.Completed);
+            var host = CreateClientConnectionHost(
+                Task.FromResult(connection),
+                new TaskCompletionSource<IClientConnection>().Task);
             var listener = new TestableDiagnosticListener();
-            var pipeName = Guid.NewGuid().ToString();
-            var dispatcherTask = Task.Run(() =>
-            {
-                var dispatcher = new ServerDispatcher(CreateNopCompilerServerHost().Object, CreateNopRequestHandler().Object, listener);
+            var keepAlive = TimeSpan.FromSeconds(1);
+            var dispatcher = new ServerDispatcher(host, listener);
             dispatcher.ListenAndDispatchConnections(keepAlive);
-            });
-
-            await RunCSharpCompile(pipeName, HelloWorldSourceText).ConfigureAwait(false);
-            await dispatcherTask.ConfigureAwait(false);
 
             Assert.Equal(1, listener.ProcessedCount);
             Assert.True(listener.LastProcessedTime.HasValue);
@@ -307,28 +209,25 @@ public async Task KeepAliveAfterSingleConnection()
         /// <summary>
         /// Ensure server respects keep alive and shuts down after processing multiple connections.
         /// </summary>
-        [Fact(Skip = "https://github.com/dotnet/roslyn/issues/4301")]
-        public async Task KeepAliveAfterMultipleConnection()
-        {
-            var keepAlive = TimeSpan.FromSeconds(1);
-            var listener = new TestableDiagnosticListener();
-            var pipeName = Guid.NewGuid().ToString();
-            var dispatcherTask = Task.Run(() =>
+        [Fact]
+        public void KeepAliveAfterMultipleConnection()
         {
-                var dispatcher = new ServerDispatcher(
-                    CreateNopCompilerServerHost().Object,
-                    new CompilerRequestHandler(CreateNopCompilerServerHost().Object),
-                    listener);
-                dispatcher.ListenAndDispatchConnections(keepAlive);
-            });
-
-            for (int i = 0; i < 5; i++)
+            var count = 5;
+            var list = new List<Task<IClientConnection>>();
+            for (var i = 0; i < count; i++)
             {
-                await RunCSharpCompile(pipeName, HelloWorldSourceText).ConfigureAwait(false);
+                var connection = CreateClientConnection(CompletionReason.Completed);
+                list.Add(Task.FromResult(connection));
             }
 
-            await dispatcherTask.ConfigureAwait(false);
-            Assert.Equal(5, listener.ProcessedCount);
+            list.Add(new TaskCompletionSource<IClientConnection>().Task);
+            var host = CreateClientConnectionHost(list.ToArray());
+            var listener = new TestableDiagnosticListener();
+            var keepAlive = TimeSpan.FromSeconds(1);
+            var dispatcher = new ServerDispatcher(host, listener);
+            dispatcher.ListenAndDispatchConnections(keepAlive);
+
+            Assert.Equal(count, listener.ProcessedCount);
             Assert.True(listener.LastProcessedTime.HasValue);
             Assert.True((DateTime.Now - listener.LastProcessedTime.Value) > keepAlive);
         }
@@ -336,70 +235,48 @@ public async Task KeepAliveAfterMultipleConnection()
         /// <summary>
         /// Ensure server respects keep alive and shuts down after processing simultaneous connections.
         /// </summary>
-        [Fact(Skip = "https://github.com/dotnet/roslyn/issues/4301")]
+        [Fact]
         public async Task KeepAliveAfterSimultaneousConnection()
         {
+            var totalCount = 2;
+            var readySource = new TaskCompletionSource<bool>();
+            var list = new List<TaskCompletionSource<ConnectionData>>();
+            var host = new Mock<IClientConnectionHost>();
+            host
+                .Setup(x => x.CreateListenTask(It.IsAny<CancellationToken>()))
+                .Returns((CancellationToken ct) =>
+                {
+                    if (list.Count < totalCount)
+                    {
+                        var source = new TaskCompletionSource<ConnectionData>();
+                        var client = CreateClientConnection(source.Task);
+                        list.Add(source);
+                        return Task.FromResult(client);
+                    }
+
+                    readySource.SetResult(true);
+                    return new TaskCompletionSource<IClientConnection>().Task;
+                });
+
             var keepAlive = TimeSpan.FromSeconds(1);
             var listener = new TestableDiagnosticListener();
-            var pipeName = Guid.NewGuid().ToString();
             var dispatcherTask = Task.Run(() =>
             {
-                var dispatcher = new ServerDispatcher(
-                    CreateNopCompilerServerHost().Object,
-                    new CompilerRequestHandler(CreateNopCompilerServerHost().Object),
-                    listener);
+                var dispatcher = new ServerDispatcher(host.Object, listener);
                 dispatcher.ListenAndDispatchConnections(keepAlive);
             });
 
-            var list = new List<Task>();
-            for (int i = 0; i < 5; i++)
-            {
-                var task = Task.Run(() => RunCSharpCompile(pipeName, HelloWorldSourceText));
-                list.Add(task);
-            }
 
-            foreach (var current in list)
+            await readySource.Task.ConfigureAwait(true);
+            foreach (var source in list)
             {
-                await current.ConfigureAwait(false);
+                source.SetResult(new ConnectionData(CompletionReason.Completed));
             }
 
-            await dispatcherTask.ConfigureAwait(false);
-            Assert.Equal(5, listener.ProcessedCount);
+            await dispatcherTask.ConfigureAwait(true);
+            Assert.Equal(totalCount, listener.ProcessedCount);
             Assert.True(listener.LastProcessedTime.HasValue);
             Assert.True((DateTime.Now - listener.LastProcessedTime.Value) > keepAlive);
         }
-
-        [Fact(Skip = "https://github.com/dotnet/roslyn/issues/4301")]
-        public async Task FirstClientCanOverrideDefaultTimeout()
-        {
-            var cts = new CancellationTokenSource();
-            var listener = new TestableDiagnosticListener();
-            TimeSpan? newTimeSpan = null;
-            var connectionSource = new TaskCompletionSource<int>();
-            var diagnosticListener = new Mock<IDiagnosticListener>();
-            diagnosticListener
-                .Setup(x => x.UpdateKeepAlive(It.IsAny<TimeSpan>()))
-                .Callback<TimeSpan>(ts => { newTimeSpan = ts; });
-            diagnosticListener
-                .Setup(x => x.ConnectionProcessed(It.IsAny<int>()))
-                .Callback<int>(count => connectionSource.SetResult(count));
-
-            var pipeName = Guid.NewGuid().ToString();
-            var dispatcherTask = Task.Run(() =>
-            {
-                var dispatcher = new ServerDispatcher(CreateNopCompilerServerHost().Object, CreateNopRequestHandler().Object, diagnosticListener.Object);
-                dispatcher.ListenAndDispatchConnections(TimeSpan.FromSeconds(1), cancellationToken: cts.Token);
-            });
-
-            var seconds = 10;
-            var response = await RunCSharpCompile(pipeName, HelloWorldSourceText, TimeSpan.FromSeconds(seconds)).ConfigureAwait(false);
-            Assert.Equal(BuildResponse.ResponseType.Completed, response.Type);
-            Assert.Equal(1, await connectionSource.Task.ConfigureAwait(false));
-            Assert.True(newTimeSpan.HasValue);
-            Assert.Equal(seconds, newTimeSpan.Value.TotalSeconds);
-
-            cts.Cancel();
-            await dispatcherTask.ConfigureAwait(false);
-        }
     }
 }

src/Compilers/Server/VBCSCompilerTests/VBCSCompilerTests.csproj

@@ -102,6 +102,7 @@
   <ItemGroup>
     <Compile Include="AnalyzerConsistencyCheckerTests.cs" />
     <Compile Include="BuildProtocolTest.cs" />
+    <Compile Include="ClientConnectionTests.cs" />
     <Compile Include="CompilerServerApiTest.cs" />
     <Compile Include="CompilerServerTests.cs" />
     <Compile Include="EndToEndDeterminismTest.cs" />

src/Compilers/VisualBasic/Test/CommandLine/TouchedFileLoggingTests.vb

@@ -169,7 +169,7 @@ End Class
 
                 Dim outWriter = New StringWriter()
                 Dim cmd = New VisualBasicCompilerServer(
-                    New DesktopCompilerServerHost(Guid.NewGuid().ToString()),
+                    New DesktopCompilerServerHost(),
                     {"/nologo",
                      "/touchedfiles:" + touchedBase,
                      source1},