first refactoring to add connection pool (#24751)

* first refactoring to add connection pool

* pool added

* removed usage of KeepAliveSession which solely used as performance benefit favoring perf over complexity.

now with pool, there is no perf advantage by KeepAliveSession. it should be only used when one wants to make connection stateful.

* change remote call signature since solution is now set automatically

* fixed test failures

* made some constant an options. and made pool off by default

handle pool shutdown better

* now turn it on by default

* PR feedbacks

* address more PR feedbacks

* more comments

* renamed ConnectionPool to ConnectionManager. people seems confused by name too much. pool is just one of thing it does, it owns connection management in general.

* PR feedbacks

* change Contract.Fail to Contract.Require so that it doesn't crash VS when OOP is killed explicitly by users.

* PR feedbacks

* address feedbacks

* put finalizer under debug flag

src/Features/Core/Portable/DesignerAttributes/AbstractDesignerAttributeService.cs

@@ -17,13 +17,9 @@ internal abstract class AbstractDesignerAttributeService : IDesignerAttributeSer
         // on remote host, so we need to make sure given input always belong to right workspace where
         // the session belong to.
         private readonly Workspace _workspace;
-        private readonly SemaphoreSlim _gate;
-
-        private KeepAliveSession _sessionDoNotAccessDirectly;
 
         protected AbstractDesignerAttributeService(Workspace workspace)
         {
-            _gate = new SemaphoreSlim(initialCount: 1);
             _workspace = workspace;
         }
 
@@ -51,33 +47,13 @@ public async Task<DesignerAttributeResult> ScanDesignerAttributesAsync(Document
             return await ScanDesignerAttributesInCurrentProcessAsync(document, cancellationToken).ConfigureAwait(false);
         }
 
-        private async Task<KeepAliveSession> TryGetKeepAliveSessionAsync(RemoteHostClient client, CancellationToken cancellationToken)
-        {
-            using (await _gate.DisposableWaitAsync(cancellationToken).ConfigureAwait(false))
-            {
-                if (_sessionDoNotAccessDirectly == null)
-                {
-                    _sessionDoNotAccessDirectly = await client.TryCreateCodeAnalysisKeepAliveSessionAsync(cancellationToken).ConfigureAwait(false);
-                }
-
-                return _sessionDoNotAccessDirectly;
-            }
-        }
-
         private async Task<DesignerAttributeResult> ScanDesignerAttributesInRemoteHostAsync(RemoteHostClient client, Document document, CancellationToken cancellationToken)
         {
-            var keepAliveSession = await TryGetKeepAliveSessionAsync(client, cancellationToken).ConfigureAwait(false);
-            if (keepAliveSession == null)
-            {
-                // The client is not currently running, so we don't know the state of the DesignerAttribute.
-                return new DesignerAttributeResult(designerAttributeArgument: null, containsErrors: false, applicable: false);
-            }
-
-            var result = await keepAliveSession.TryInvokeAsync<DesignerAttributeResult>(
+            return await client.TryRunCodeAnalysisRemoteAsync<DesignerAttributeResult>(
+                document.Project.Solution,
                 nameof(IRemoteDesignerAttributeService.ScanDesignerAttributesAsync),
-                document.Project.Solution, new object[] { document.Id }, cancellationToken).ConfigureAwait(false);
-
-            return result;
+                document.Id,
+                cancellationToken).ConfigureAwait(false);
         }
 
         private async Task<DesignerAttributeResult> ScanDesignerAttributesInCurrentProcessAsync(Document document, CancellationToken cancellationToken)

src/Features/Core/Portable/DesignerAttributes/IRemoteDesignerAttributeService.cs

@@ -8,6 +8,6 @@ namespace Microsoft.CodeAnalysis.DesignerAttributes
 {
     internal interface IRemoteDesignerAttributeService
     {
-        Task<DesignerAttributeResult> ScanDesignerAttributesAsync(PinnedSolutionInfo solutionInfo, DocumentId documentId, CancellationToken cancellationToken);
+        Task<DesignerAttributeResult> ScanDesignerAttributesAsync(DocumentId documentId, CancellationToken cancellationToken);
     }
 }

src/Features/Core/Portable/TodoComments/AbstractTodoCommentService.cs

@@ -19,13 +19,9 @@ internal abstract class AbstractTodoCommentService : ITodoCommentService
         // on remote host, so we need to make sure given input always belong to right workspace where
         // the session belong to.
         private readonly Workspace _workspace;
-        private readonly SemaphoreSlim _gate;
-
-        private KeepAliveSession _sessionDoNotAccessDirectly;
 
         protected AbstractTodoCommentService(Workspace workspace)
         {
-            _gate = new SemaphoreSlim(initialCount: 1);
             _workspace = workspace;
         }
 
@@ -61,34 +57,14 @@ public async Task<IList<TodoComment>> GetTodoCommentsAsync(Document document, IL
         private async Task<IList<TodoComment>> GetTodoCommentsInRemoteHostAsync(
             RemoteHostClient client, Document document, IList<TodoCommentDescriptor> commentDescriptors, CancellationToken cancellationToken)
         {
-            var keepAliveSession = await TryGetKeepAliveSessionAsync(client, cancellationToken).ConfigureAwait(false);
-            if (keepAliveSession == null)
-            {
-                // The client is not currently running, so we don't have any results.
-                return SpecializedCollections.EmptyList<TodoComment>();
-            }
-
-            var result = await keepAliveSession.TryInvokeAsync<IList<TodoComment>>(
-                nameof(IRemoteTodoCommentService.GetTodoCommentsAsync),
+            var result = await client.TryRunCodeAnalysisRemoteAsync<IList<TodoComment>>(
                 document.Project.Solution,
+                nameof(IRemoteTodoCommentService.GetTodoCommentsAsync),
                 new object[] { document.Id, commentDescriptors }, cancellationToken).ConfigureAwait(false);
 
             return result ?? SpecializedCollections.EmptyList<TodoComment>();
         }
 
-        private async Task<KeepAliveSession> TryGetKeepAliveSessionAsync(RemoteHostClient client, CancellationToken cancellationToken)
-        {
-            using (await _gate.DisposableWaitAsync(cancellationToken).ConfigureAwait(false))
-            {
-                if (_sessionDoNotAccessDirectly == null)
-                {
-                    _sessionDoNotAccessDirectly = await client.TryCreateCodeAnalysisKeepAliveSessionAsync(cancellationToken).ConfigureAwait(false);
-                }
-
-                return _sessionDoNotAccessDirectly;
-            }
-        }
-
         private async Task<IList<TodoComment>> GetTodoCommentsInCurrentProcessAsync(
             Document document, IList<TodoCommentDescriptor> commentDescriptors, CancellationToken cancellationToken)
         {

src/Features/Core/Portable/TodoComments/IRemoteTodoCommentService.cs

@@ -13,6 +13,6 @@ namespace Microsoft.CodeAnalysis.TodoComments
     /// </summary>
     internal interface IRemoteTodoCommentService
     {
-        Task<IList<TodoComment>> GetTodoCommentsAsync(PinnedSolutionInfo solutionInfo, DocumentId documentId, IList<TodoCommentDescriptor> commentDescriptors, CancellationToken cancellationToken);
+        Task<IList<TodoComment>> GetTodoCommentsAsync(DocumentId documentId, IList<TodoCommentDescriptor> commentDescriptors, CancellationToken cancellationToken);
     }
 }

src/VisualStudio/Core/Def/Implementation/Remote/JsonRpcConnection.cs

@@ -6,7 +6,6 @@
 using System.IO;
 using System.Threading;
 using System.Threading.Tasks;
-using Microsoft.CodeAnalysis.Execution;
 using Microsoft.CodeAnalysis.Remote;
 using Roslyn.Utilities;
 
@@ -32,11 +31,6 @@ internal class JsonRpcConnection : RemoteHostClient.Connection
             _remoteDataRpc = dataRpc;
         }
 
-        protected override async Task OnRegisterPinnedRemotableDataScopeAsync(PinnedRemotableDataScope scope)
-        {
-            await InvokeAsync(WellKnownServiceHubServices.ServiceHubServiceBase_Initialize, new object[] { scope.SolutionInfo }, CancellationToken.None).ConfigureAwait(false);
-        }
-
         public override Task InvokeAsync(string targetName, IReadOnlyList<object> arguments, CancellationToken cancellationToken)
         {
             return _serviceRpc.InvokeAsync(targetName, arguments, cancellationToken);
@@ -57,13 +51,16 @@ public override Task<T> InvokeAsync<T>(string targetName, IReadOnlyList<object>
             return _serviceRpc.InvokeAsync<T>(targetName, arguments, funcWithDirectStreamAsync, cancellationToken);
         }
 
-        protected override void OnDisposed()
+        protected override void Dispose(bool disposing)
         {
-            base.OnDisposed();
+            if (disposing)
+            {
+                // dispose service and snapshot channels
+                _serviceRpc.Dispose();
+                _remoteDataRpc.Dispose();
+            }
 
-            // dispose service and snapshot channels
-            _serviceRpc.Dispose();
-            _remoteDataRpc.Dispose();
+            base.Dispose(disposing);
         }
 
         /// <summary>

src/VisualStudio/Core/Def/Implementation/Remote/RemoteHostOptions.cs

@@ -53,6 +53,17 @@ internal static class RemoteHostOptions
             storageLocations: new LocalUserProfileStorageLocation(InternalFeatureOnOffOptions.LocalRegistryPath + nameof(OOP64Bit)));
 
         public static readonly Option<bool> RemoteHostTest = new Option<bool>(nameof(InternalFeatureOnOffOptions), nameof(RemoteHostTest), defaultValue: false);
+
+        public static readonly Option<bool> EnableConnectionPool = new Option<bool>(
+            nameof(InternalFeatureOnOffOptions), nameof(EnableConnectionPool), defaultValue: true,
+            storageLocations: new LocalUserProfileStorageLocation(InternalFeatureOnOffOptions.LocalRegistryPath + nameof(EnableConnectionPool)));
+
+        /// <summary>
+        /// default 15 is chosen which is big enough but not too big for service hub to handle
+        /// </summary>
+        public static readonly Option<int> MaxPoolConnection = new Option<int>(
+            nameof(InternalFeatureOnOffOptions), nameof(MaxPoolConnection), defaultValue: 15,
+            storageLocations: new LocalUserProfileStorageLocation(InternalFeatureOnOffOptions.LocalRegistryPath + nameof(MaxPoolConnection)));
     }
 
     [ExportOptionProvider, Shared]
@@ -64,6 +75,8 @@ internal class RemoteHostOptionsProvider : IOptionProvider
             RemoteHostOptions.RequestServiceTimeoutInMS,
             RemoteHostOptions.RestartRemoteHostAllowed,
             RemoteHostOptions.OOP64Bit,
-            RemoteHostOptions.RemoteHostTest);
+            RemoteHostOptions.RemoteHostTest,
+            RemoteHostOptions.EnableConnectionPool,
+            RemoteHostOptions.MaxPoolConnection);
     }
 }

src/VisualStudio/Core/Def/Implementation/Remote/ServiceHubRemoteHostClient.ConnectionManager.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.Concurrent;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.CodeAnalysis.Remote;
+using Microsoft.ServiceHub.Client;
+using Roslyn.Utilities;
+
+namespace Microsoft.VisualStudio.LanguageServices.Remote
+{
+    internal sealed partial class ServiceHubRemoteHostClient : RemoteHostClient
+    {
+        private partial class ConnectionManager
+        {
+            private readonly HubClient _hubClient;
+            private readonly HostGroup _hostGroup;
+            private readonly TimeSpan _timeout;
+
+            private readonly ReaderWriterLockSlim _shutdownLock;
+            private readonly ReferenceCountedDisposable<RemotableDataJsonRpc> _remotableDataRpc;
+
+            private readonly int _maxPoolConnections;
+
+            // keyed to serviceName. each connection is for specific service such as CodeAnalysisService
+            private readonly ConcurrentDictionary<string, ConcurrentQueue<JsonRpcConnection>> _pools;
+
+            // indicate whether pool should be used.
+            private readonly bool _enableConnectionPool;
+
+            public ConnectionManager(
+                HubClient hubClient,
+                HostGroup hostGroup,
+                bool enableConnectionPool,
+                int maxPoolConnection,
+                TimeSpan timeout,
+                ReferenceCountedDisposable<RemotableDataJsonRpc> remotableDataRpc)
+            {
+                _hubClient = hubClient;
+                _hostGroup = hostGroup;
+                _timeout = timeout;
+
+                _remotableDataRpc = remotableDataRpc;
+                _maxPoolConnections = maxPoolConnection;
+
+                // initial value 4 is chosen to stop concurrent dictionary creating too many locks.
+                // and big enough for all our services such as codeanalysis, remotehost, snapshot and etc services
+                _pools = new ConcurrentDictionary<string, ConcurrentQueue<JsonRpcConnection>>(concurrencyLevel: 4, capacity: 4);
+
+                _enableConnectionPool = enableConnectionPool;
+                _shutdownLock = new ReaderWriterLockSlim(LockRecursionPolicy.NoRecursion);
+            }
+
+            public Task<Connection> TryCreateConnectionAsync(string serviceName, object callbackTarget, CancellationToken cancellationToken)
+            {
+                // pool is not enabled by option
+                if (!_enableConnectionPool)
+                {
+                    // RemoteHost is allowed to be restarted by IRemoteHostClientService.RequestNewRemoteHostAsync
+                    // when that happens, existing Connection will keep working until they get disposed.
+                    //
+                    // now question is when someone calls RemoteHostClient.TryGetConnection for the client that got
+                    // shutdown, whether it should gracefully handle that request or fail after shutdown.
+                    // for current expected usage case where new remoteHost is only created when new solution is added,
+                    // we should be fine on failing after shutdown.
+                    //
+                    // but, at some point, if we want to support RemoteHost being restarted at any random point, 
+                    // we need to revisit this to support such case by creating new temporary connections.
+                    // for now, I dropped it since it felt over-designing when there is no usage case for that yet.
+                    return TryCreateNewConnectionAsync(serviceName, callbackTarget, cancellationToken);
+                }
+
+                // when callbackTarget is given, we can't share/pool connection since callbackTarget attaches a state to connection.
+                // so connection is only valid for that specific callbackTarget. it is up to the caller to keep connection open
+                // if he wants to reuse same connection
+                if (callbackTarget != null)
+                {
+                    return TryCreateNewConnectionAsync(serviceName, callbackTarget, cancellationToken);
+                }
+
+                return TryGetConnectionFromPoolAsync(serviceName, cancellationToken);
+            }
+
+            private async Task<Connection> TryGetConnectionFromPoolAsync(string serviceName, CancellationToken cancellationToken)
+            {
+                var queue = _pools.GetOrAdd(serviceName, _ => new ConcurrentQueue<JsonRpcConnection>());
+                if (queue.TryDequeue(out var connection))
+                {
+                    return new PooledConnection(this, serviceName, connection);
+                }
+
+                var newConnection = (JsonRpcConnection)await TryCreateNewConnectionAsync(serviceName, callbackTarget: null, cancellationToken).ConfigureAwait(false);
+                if (newConnection == null)
+                {
+                    // we might not get new connection if we are either shutdown explicitly or due to OOP terminated
+                    return null;
+                }
+
+                return new PooledConnection(this, serviceName, newConnection);
+            }
+
+            private async Task<Connection> TryCreateNewConnectionAsync(string serviceName, object callbackTarget, CancellationToken cancellationToken)
+            {
+                var dataRpc = _remotableDataRpc.TryAddReference();
+                if (dataRpc == null)
+                {
+                    // dataRpc is disposed. this can happen if someone killed remote host process while there is
+                    // no other one holding the data connection.
+                    // in those error case, don't crash but return null. this method is TryCreate since caller expects it to return null
+                    // on such error situation.
+                    return null;
+                }
+
+                // get stream from service hub to communicate service specific information
+                // this is what consumer actually use to communicate information
+                var serviceStream = await Connections.RequestServiceAsync(_hubClient, serviceName, _hostGroup, _timeout, cancellationToken).ConfigureAwait(false);
+
+                return new JsonRpcConnection(_hubClient.Logger, callbackTarget, serviceStream, dataRpc);
+            }
+
+            private void Free(string serviceName, JsonRpcConnection connection)
+            {
+                using (_shutdownLock.DisposableRead())
+                {
+                    if (!_enableConnectionPool)
+                    {
+                        // pool is not being used.
+                        connection.Dispose();
+                        return;
+                    }
+
+                    // queue must exist
+                    var queue = _pools[serviceName];
+                    if (queue.Count >= _maxPoolConnections)
+                    {
+                        // let the connection actually go away
+                        connection.Dispose();
+                        return;
+                    }
+
+                    // pool the connection
+                    queue.Enqueue(connection);
+                }
+            }
+
+            public void Shutdown()
+            {
+                using (_shutdownLock.DisposableWrite())
+                {
+                    // let ref count this one is holding go
+                    _remotableDataRpc.Dispose();
+
+                    // let all connections in the pool to go away
+                    foreach (var (_, queue) in _pools)
+                    {
+                        while (queue.TryDequeue(out var connection))
+                        {
+                            connection.Dispose();
+                        }
+                    }
+
+                    _pools.Clear();
+                }
+            }
+        }
+    }
+}

src/VisualStudio/Core/Def/Implementation/Remote/ServiceHubRemoteHostClient.Connections.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.Diagnostics;
+using System.IO;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.CodeAnalysis;
+using Microsoft.CodeAnalysis.ErrorReporting;
+using Microsoft.CodeAnalysis.Extensions;
+using Microsoft.CodeAnalysis.Remote;
+using Microsoft.ServiceHub.Client;
+using Roslyn.Utilities;
+using StreamJsonRpc;
+
+namespace Microsoft.VisualStudio.LanguageServices.Remote
+{
+    internal sealed partial class ServiceHubRemoteHostClient : RemoteHostClient
+    {
+        private static class Connections
+        {
+            /// <summary>
+            /// call <paramref name="funcAsync"/> and retry up to <paramref name="timeout"/> if the call throws
+            /// <typeparamref name="TException"/>. any other exception from the call won't be handled here.
+            /// </summary>
+            public static async Task<TResult> RetryRemoteCallAsync<TException, TResult>(
+                Func<Task<TResult>> funcAsync,
+                TimeSpan timeout,
+                CancellationToken cancellationToken) where TException : Exception
+            {
+                const int retry_delayInMS = 50;
+
+                using (var pooledStopwatch = SharedPools.Default<Stopwatch>().GetPooledObject())
+                {
+                    var watch = pooledStopwatch.Object;
+                    watch.Start();
+
+                    while (watch.Elapsed < timeout)
+                    {
+                        cancellationToken.ThrowIfCancellationRequested();
+
+                        try
+                        {
+                            return await funcAsync().ConfigureAwait(false);
+                        }
+                        catch (TException)
+                        {
+                            // throw cancellation token if operation is cancelled
+                            cancellationToken.ThrowIfCancellationRequested();
+                        }
+
+                        // wait for retry_delayInMS before next try
+                        await Task.Delay(retry_delayInMS, cancellationToken).ConfigureAwait(false);
+
+                        ReportTimeout(watch);
+                    }
+                }
+
+                // operation timed out, more than we are willing to wait
+                ShowInfoBar();
+
+                // user didn't ask for cancellation, but we can't fullfill this request. so we
+                // create our own cancellation token and then throw it. this doesn't guarantee
+                // 100% that we won't crash, but this is at least safest way we know until user
+                // restart VS (with info bar)
+                using (var ownCancellationSource = new CancellationTokenSource())
+                {
+                    ownCancellationSource.Cancel();
+                    ownCancellationSource.Token.ThrowIfCancellationRequested();
+                }
+
+                throw ExceptionUtilities.Unreachable;
+            }
+
+            public static async Task<Stream> RequestServiceAsync(
+                HubClient client,
+                string serviceName,
+                HostGroup hostGroup,
+                TimeSpan timeout,
+                CancellationToken cancellationToken)
+            {
+                const int max_retry = 10;
+                const int retry_delayInMS = 50;
+
+                RemoteInvocationException lastException = null;
+
+                var descriptor = new ServiceDescriptor(serviceName) { HostGroup = hostGroup };
+
+                // call to get service can fail due to this bug - devdiv#288961 or more.
+                // until root cause is fixed, we decide to have retry rather than fail right away
+                for (var i = 0; i < max_retry; i++)
+                {
+                    try
+                    {
+                        // we are wrapping HubClient.RequestServiceAsync since we can't control its internal timeout value ourselves.
+                        // we have bug opened to track the issue.
+                        // https://devdiv.visualstudio.com/DefaultCollection/DevDiv/Editor/_workitems?id=378757&fullScreen=false&_a=edit
+
+                        // retry on cancellation token since HubClient will throw its own cancellation token
+                        // when it couldn't connect to service hub service for some reasons
+                        // (ex, OOP process GC blocked and not responding to request)
+                        //
+                        // we have double re-try here. we have these 2 seperated since 2 retries are for different problems.
+                        // as noted by 2 different issues above at the start of each 2 different retries.
+                        // first retry most likely deal with real issue on servicehub, second retry (cancellation) is to deal with
+                        // by design servicehub behavior we don't want to use.
+                        return await RetryRemoteCallAsync<OperationCanceledException, Stream>(
+                            () => client.RequestServiceAsync(descriptor, cancellationToken),
+                            timeout,
+                            cancellationToken).ConfigureAwait(false);
+                    }
+                    catch (RemoteInvocationException ex)
+                    {
+                        // save info only if it failed with different issue than before.
+                        if (lastException?.Message != ex.Message)
+                        {
+                            // RequestServiceAsync should never fail unless service itself is actually broken.
+                            // So far, we catched multiple issues from this NFW. so we will keep this NFW.
+                            ex.ReportServiceHubNFW("RequestServiceAsync Failed");
+
+                            lastException = ex;
+                        }
+                    }
+
+                    // wait for retry_delayInMS before next try
+                    await Task.Delay(retry_delayInMS, cancellationToken).ConfigureAwait(false);
+                }
+
+                // crash right away to get better dump. otherwise, we will get dump from async exception
+                // which most likely lost all valuable data
+                FatalError.ReportUnlessCanceled(lastException);
+                GC.KeepAlive(lastException);
+
+                // unreachable
+                throw ExceptionUtilities.Unreachable;
+            }
+
+            #region code related to make diagnosis easier later
+
+            private static readonly TimeSpan s_reportTimeout = TimeSpan.FromMinutes(10);
+            private static bool s_timeoutReported = false;
+
+            private static void ReportTimeout(Stopwatch watch)
+            {
+                // if we tried for 10 min and still couldn't connect. NFW (non fatal watson) some data
+                if (!s_timeoutReported && watch.Elapsed > s_reportTimeout)
+                {
+                    s_timeoutReported = true;
+
+                    // report service hub logs along with dump
+                    (new Exception("RequestServiceAsync Timeout")).ReportServiceHubNFW("RequestServiceAsync Timeout");
+                }
+            }
+
+            private static bool s_infoBarReported = false;
+
+            private static void ShowInfoBar()
+            {
+                // use info bar to show warning to users
+                if (CodeAnalysis.PrimaryWorkspace.Workspace != null && !s_infoBarReported)
+                {
+                    // do not report it multiple times
+                    s_infoBarReported = true;
+
+                    // use info bar to show warning to users
+                    CodeAnalysis.PrimaryWorkspace.Workspace.Services.GetService<IErrorReportingService>()?.ShowGlobalErrorInfo(
+                        ServicesVSResources.Unfortunately_a_process_used_by_Visual_Studio_has_encountered_an_unrecoverable_error_We_recommend_saving_your_work_and_then_closing_and_restarting_Visual_Studio);
+                }
+            }
+            #endregion
+        }
+    }
+}

src/VisualStudio/Core/Def/Implementation/Remote/ServiceHubRemoteHostClient.PooledConnection.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.IO;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.CodeAnalysis.Remote;
+
+namespace Microsoft.VisualStudio.LanguageServices.Remote
+{
+    internal sealed partial class ServiceHubRemoteHostClient : RemoteHostClient
+    {
+        private partial class ConnectionManager
+        {
+            private class PooledConnection : Connection
+            {
+                private readonly ConnectionManager _connectionManager;
+                private readonly string _serviceName;
+                private readonly JsonRpcConnection _connection;
+
+                public PooledConnection(ConnectionManager pools, string serviceName, JsonRpcConnection connection)
+                {
+                    _connectionManager = pools;
+                    _serviceName = serviceName;
+                    _connection = connection;
+                }
+
+                public override Task InvokeAsync(string targetName, IReadOnlyList<object> arguments, CancellationToken cancellationToken) =>
+                    _connection.InvokeAsync(targetName, arguments, cancellationToken);
+
+                public override Task<T> InvokeAsync<T>(string targetName, IReadOnlyList<object> arguments, CancellationToken cancellationToken) =>
+                    _connection.InvokeAsync<T>(targetName, arguments, cancellationToken);
+
+                public override Task InvokeAsync(
+                    string targetName, IReadOnlyList<object> arguments,
+                    Func<Stream, CancellationToken, Task> funcWithDirectStreamAsync, CancellationToken cancellationToken) =>
+                    _connection.InvokeAsync(targetName, arguments, funcWithDirectStreamAsync, cancellationToken);
+
+                public override Task<T> InvokeAsync<T>(
+                    string targetName, IReadOnlyList<object> arguments,
+                    Func<Stream, CancellationToken, Task<T>> funcWithDirectStreamAsync, CancellationToken cancellationToken) =>
+                    _connection.InvokeAsync<T>(targetName, arguments, funcWithDirectStreamAsync, cancellationToken);
+
+                protected override void Dispose(bool disposing)
+                {
+                    if (disposing)
+                    {
+                        _connectionManager.Free(_serviceName, _connection);
+                    }
+
+                    base.Dispose(disposing);
+                }
+            }
+        }
+    }
+}

src/VisualStudio/Core/Def/Implementation/Remote/ServiceHubRemoteHostClient.cs

@@ -6,11 +6,7 @@
 using System.IO;
 using System.Threading;
 using System.Threading.Tasks;
-using Microsoft.CodeAnalysis;
 using Microsoft.CodeAnalysis.Editor.Shared.Utilities;
-using Microsoft.CodeAnalysis.ErrorReporting;
-using Microsoft.CodeAnalysis.Execution;
-using Microsoft.CodeAnalysis.Extensions;
 using Microsoft.CodeAnalysis.Internal.Log;
 using Microsoft.CodeAnalysis.Notification;
 using Microsoft.CodeAnalysis.Remote;
@@ -35,12 +31,8 @@ private enum GlobalNotificationState
 
         private static int s_instanceId = 0;
 
-        private readonly HubClient _hubClient;
-        private readonly HostGroup _hostGroup;
-        private readonly TimeSpan _timeout;
-
         private readonly JsonRpc _rpc;
-        private readonly ReferenceCountedDisposable<RemotableDataJsonRpc> _remotableDataRpc;
+        private readonly ConnectionManager _connectionManager;
 
         /// <summary>
         /// Lock for the <see cref="_globalNotificationsTask"/> task chain.  Each time we hear 
@@ -63,7 +55,7 @@ private enum GlobalNotificationState
                 // Retry (with timeout) until we can connect to RemoteHost (service hub process). 
                 // we are seeing cases where we failed to connect to service hub process when a machine is under heavy load.
                 // (see https://devdiv.visualstudio.com/DevDiv/_workitems/edit/481103 as one of example)
-                var instance = await RetryRemoteCallAsync<IOException, ServiceHubRemoteHostClient>(
+                var instance = await Connections.RetryRemoteCallAsync<IOException, ServiceHubRemoteHostClient>(
                     () => CreateWorkerAsync(workspace, primary, timeout, cancellationToken), timeout, cancellationToken).ConfigureAwait(false);
 
                 instance.Started();
@@ -88,11 +80,18 @@ public static async Task<ServiceHubRemoteHostClient> CreateWorkerAsync(Workspace
                 var current = $"VS ({Process.GetCurrentProcess().Id}) ({currentInstanceId})";
 
                 var hostGroup = new HostGroup(current);
-                var remoteHostStream = await RequestServiceAsync(primary, WellKnownRemoteHostServices.RemoteHostService, hostGroup, timeout, cancellationToken).ConfigureAwait(false);
+                var remoteHostStream = await Connections.RequestServiceAsync(primary, WellKnownRemoteHostServices.RemoteHostService, hostGroup, timeout, cancellationToken).ConfigureAwait(false);
 
                 var remotableDataRpc = new RemotableDataJsonRpc(
-                    workspace, primary.Logger, await RequestServiceAsync(primary, WellKnownServiceHubServices.SnapshotService, hostGroup, timeout, cancellationToken).ConfigureAwait(false));
-                client = new ServiceHubRemoteHostClient(workspace, primary, hostGroup, new ReferenceCountedDisposable<RemotableDataJsonRpc>(remotableDataRpc), remoteHostStream);
+                                          workspace, primary.Logger,
+                                          await Connections.RequestServiceAsync(primary, WellKnownServiceHubServices.SnapshotService, hostGroup, timeout, cancellationToken).ConfigureAwait(false));
+
+                var enableConnectionPool = workspace.Options.GetOption(RemoteHostOptions.EnableConnectionPool);
+                var maxConnection = workspace.Options.GetOption(RemoteHostOptions.MaxPoolConnection);
+
+                var connectionManager = new ConnectionManager(primary, hostGroup, enableConnectionPool, maxConnection, timeout, new ReferenceCountedDisposable<RemotableDataJsonRpc>(remotableDataRpc));
+
+                client = new ServiceHubRemoteHostClient(workspace, connectionManager, remoteHostStream);
 
                 var uiCultureLCID = CultureInfo.CurrentUICulture.LCID;
                 var cultureLCID = CultureInfo.CurrentCulture.LCID;
@@ -147,18 +146,11 @@ internal static async Task RegisterWorkspaceHostAsync(Workspace workspace, Remot
 
         private ServiceHubRemoteHostClient(
             Workspace workspace,
-            HubClient hubClient,
-            HostGroup hostGroup,
-            ReferenceCountedDisposable<RemotableDataJsonRpc> remotableDataRpc,
+            ConnectionManager connectionManager,
             Stream stream)
             : base(workspace)
         {
-            Contract.ThrowIfNull(remotableDataRpc);
-
-            _hubClient = hubClient;
-            _hostGroup = hostGroup;
-            _timeout = TimeSpan.FromMilliseconds(workspace.Options.GetOption(RemoteHostOptions.RequestServiceTimeoutInMS));
-            _remotableDataRpc = remotableDataRpc;
+            _connectionManager = connectionManager;
 
             _rpc = new JsonRpc(new JsonRpcMessageHandler(stream, stream), target: this);
             _rpc.JsonSerializer.Converters.Add(AggregateJsonConverter.Instance);
@@ -169,23 +161,9 @@ internal static async Task RegisterWorkspaceHostAsync(Workspace workspace, Remot
             _rpc.StartListening();
         }
 
-        public override async Task<Connection> TryCreateConnectionAsync(string serviceName, object callbackTarget, CancellationToken cancellationToken)
+        public override Task<Connection> TryCreateConnectionAsync(string serviceName, object callbackTarget, CancellationToken cancellationToken)
         {
-            var dataRpc = _remotableDataRpc.TryAddReference();
-            if (dataRpc == null)
-            {
-                // dataRpc is disposed. this can happen if someone killed remote host process while there is
-                // no other one holding the data connection.
-                // in those error case, don't crash but return null. this method is TryCreate since caller expects it to return null
-                // on such error situation.
-                return null;
-            }
-
-            // get stream from service hub to communicate service specific information
-            // this is what consumer actually use to communicate information
-            var serviceStream = await RequestServiceAsync(_hubClient, serviceName, _hostGroup, _timeout, cancellationToken).ConfigureAwait(false);
-
-            return new JsonRpcConnection(_hubClient.Logger, callbackTarget, serviceStream, dataRpc);
+            return _connectionManager.TryCreateConnectionAsync(serviceName, callbackTarget, cancellationToken);
         }
 
         protected override void OnStarted()
@@ -203,7 +181,7 @@ protected override void OnStopped()
             UnregisterGlobalOperationNotifications();
             _rpc.Disconnected -= OnRpcDisconnected;
             _rpc.Dispose();
-            _remotableDataRpc.Dispose();
+            _connectionManager.Shutdown();
         }
 
         private void RegisterGlobalOperationNotifications()
@@ -296,150 +274,5 @@ private void OnRpcDisconnected(object sender, JsonRpcDisconnectedEventArgs e)
         {
             Stopped();
         }
-
-        /// <summary>
-        /// call <paramref name="funcAsync"/> and retry up to <paramref name="timeout"/> if the call throws
-        /// <typeparamref name="TException"/>. any other exception from the call won't be handled here.
-        /// </summary>
-        private static async Task<TResult> RetryRemoteCallAsync<TException, TResult>(
-            Func<Task<TResult>> funcAsync,
-            TimeSpan timeout,
-            CancellationToken cancellationToken) where TException : Exception
-        {
-            const int retry_delayInMS = 50;
-
-            using (var pooledStopwatch = SharedPools.Default<Stopwatch>().GetPooledObject())
-            {
-                var watch = pooledStopwatch.Object;
-                watch.Start();
-
-                while (watch.Elapsed < timeout)
-                {
-                    cancellationToken.ThrowIfCancellationRequested();
-
-                    try
-                    {
-                        return await funcAsync().ConfigureAwait(false);
-                    }
-                    catch (TException)
-                    {
-                        // throw cancellation token if operation is cancelled
-                        cancellationToken.ThrowIfCancellationRequested();
-                    }
-
-                    // wait for retry_delayInMS before next try
-                    await Task.Delay(retry_delayInMS, cancellationToken).ConfigureAwait(false);
-
-                    ReportTimeout(watch);
-                }
-            }
-
-            // operation timed out, more than we are willing to wait
-            ShowInfoBar();
-
-            // user didn't ask for cancellation, but we can't fullfill this request. so we
-            // create our own cancellation token and then throw it. this doesn't guarantee
-            // 100% that we won't crash, but this is at least safest way we know until user
-            // restart VS (with info bar)
-            using (var ownCancellationSource = new CancellationTokenSource())
-            {
-                ownCancellationSource.Cancel();
-                ownCancellationSource.Token.ThrowIfCancellationRequested();
-            }
-
-            throw ExceptionUtilities.Unreachable;
-        }
-
-        private static async Task<Stream> RequestServiceAsync(
-            HubClient client,
-            string serviceName,
-            HostGroup hostGroup,
-            TimeSpan timeout,
-            CancellationToken cancellationToken = default)
-        {
-            const int max_retry = 10;
-            const int retry_delayInMS = 50;
-
-            Exception lastException = null;
-
-            var descriptor = new ServiceDescriptor(serviceName) { HostGroup = hostGroup };
-
-            // call to get service can fail due to this bug - devdiv#288961 or more.
-            // until root cause is fixed, we decide to have retry rather than fail right away
-            for (var i = 0; i < max_retry; i++)
-            {
-                try
-                {
-                    // we are wrapping HubClient.RequestServiceAsync since we can't control its internal timeout value ourselves.
-                    // we have bug opened to track the issue.
-                    // https://devdiv.visualstudio.com/DefaultCollection/DevDiv/Editor/_workitems?id=378757&fullScreen=false&_a=edit
-
-                    // retry on cancellation token since HubClient will throw its own cancellation token
-                    // when it couldn't connect to service hub service for some reasons
-                    // (ex, OOP process GC blocked and not responding to request)
-                    return await RetryRemoteCallAsync<OperationCanceledException, Stream>(
-                        () => client.RequestServiceAsync(descriptor, cancellationToken),
-                        timeout,
-                        cancellationToken).ConfigureAwait(false);
-                }
-                catch (RemoteInvocationException ex)
-                {
-                    // save info only if it failed with different issue than before.
-                    if (lastException?.Message != ex.Message)
-                    {
-                        // RequestServiceAsync should never fail unless service itself is actually broken.
-                        // So far, we catched multiple issues from this NFW. so we will keep this NFW.
-                        ex.ReportServiceHubNFW("RequestServiceAsync Failed");
-
-                        lastException = ex;
-                    }
-                }
-
-                // wait for retry_delayInMS before next try
-                await Task.Delay(retry_delayInMS, cancellationToken).ConfigureAwait(false);
-            }
-
-            // crash right away to get better dump. otherwise, we will get dump from async exception
-            // which most likely lost all valuable data
-            FatalError.ReportUnlessCanceled(lastException);
-            GC.KeepAlive(lastException);
-
-            // unreachable
-            throw ExceptionUtilities.Unreachable;
-        }
-
-        #region code related to make diagnosis easier later
-
-        private static readonly TimeSpan s_reportTimeout = TimeSpan.FromMinutes(10);
-        private static bool s_timeoutReported = false;
-
-        private static void ReportTimeout(Stopwatch watch)
-        {
-            // if we tried for 10 min and still couldn't connect. NFW (non fatal watson) some data
-            if (!s_timeoutReported && watch.Elapsed > s_reportTimeout)
-            {
-                s_timeoutReported = true;
-
-                // report service hub logs along with dump
-                (new Exception("RequestServiceAsync Timeout")).ReportServiceHubNFW("RequestServiceAsync Timeout");
-            }
-        }
-
-        private static bool s_infoBarReported = false;
-
-        private static void ShowInfoBar()
-        {
-            // use info bar to show warning to users
-            if (CodeAnalysis.PrimaryWorkspace.Workspace != null && !s_infoBarReported)
-            {
-                // do not report it multiple times
-                s_infoBarReported = true;
-
-                // use info bar to show warning to users
-                CodeAnalysis.PrimaryWorkspace.Workspace.Services.GetService<IErrorReportingService>()?.ShowGlobalErrorInfo(
-                    ServicesVSResources.Unfortunately_a_process_used_by_Visual_Studio_has_encountered_an_unrecoverable_error_We_recommend_saving_your_work_and_then_closing_and_restarting_Visual_Studio);
-            }
-        }
-        #endregion
     }
 }

src/VisualStudio/Core/Test.Next/Services/ServiceHubServicesTests.cs

@@ -77,16 +77,13 @@ public async Task TestTodoComments()
 
                 var solution = workspace.CurrentSolution;
 
-                var keepAliveSession = await client.TryCreateCodeAnalysisKeepAliveSessionAsync(CancellationToken.None);
-                var comments = await keepAliveSession.TryInvokeAsync<IList<TodoComment>>(
-                    nameof(IRemoteTodoCommentService.GetTodoCommentsAsync),
+                var comments = await client.TryRunCodeAnalysisRemoteAsync<IList<TodoComment>>(
                     solution,
+                    nameof(IRemoteTodoCommentService.GetTodoCommentsAsync),
                     new object[] { solution.Projects.First().DocumentIds.First(), ImmutableArray.Create(new TodoCommentDescriptor("TODO", 0)) },
                     CancellationToken.None);
 
                 Assert.Equal(comments.Count, 1);
-
-                keepAliveSession.Shutdown();
             }
         }
 
@@ -102,11 +99,10 @@ class Test { }";
 
                 var solution = workspace.CurrentSolution;
 
-                var keepAliveSession = await client.TryCreateCodeAnalysisKeepAliveSessionAsync(CancellationToken.None);
-                var result = await keepAliveSession.TryInvokeAsync<DesignerAttributeResult>(
-                    nameof(IRemoteDesignerAttributeService.ScanDesignerAttributesAsync),
+                var result = await client.TryRunCodeAnalysisRemoteAsync<DesignerAttributeResult>(
                     solution,
-                    new object[] { solution.Projects.First().DocumentIds.First() },
+                    nameof(IRemoteDesignerAttributeService.ScanDesignerAttributesAsync),
+                    solution.Projects.First().DocumentIds.First(),
                     CancellationToken.None);
 
                 Assert.Equal(result.DesignerAttributeArgument, "Form");

src/Workspaces/Core/Portable/Remote/RemoteHostClient.cs

@@ -5,7 +5,6 @@
 using System.IO;
 using System.Threading;
 using System.Threading.Tasks;
-using Microsoft.CodeAnalysis.Execution;
 using Roslyn.Utilities;
 
 namespace Microsoft.CodeAnalysis.Remote
@@ -105,19 +104,12 @@ protected Connection()
                 _disposed = false;
             }
 
-            protected abstract Task OnRegisterPinnedRemotableDataScopeAsync(PinnedRemotableDataScope scope);
-
-            public virtual Task RegisterPinnedRemotableDataScopeAsync(PinnedRemotableDataScope scope)
-            {
-                return OnRegisterPinnedRemotableDataScopeAsync(scope);
-            }
-
             public abstract Task InvokeAsync(string targetName, IReadOnlyList<object> arguments, CancellationToken cancellationToken);
             public abstract Task<T> InvokeAsync<T>(string targetName, IReadOnlyList<object> arguments, CancellationToken cancellationToken);
             public abstract Task InvokeAsync(string targetName, IReadOnlyList<object> arguments, Func<Stream, CancellationToken, Task> funcWithDirectStreamAsync, CancellationToken cancellationToken);
             public abstract Task<T> InvokeAsync<T>(string targetName, IReadOnlyList<object> arguments, Func<Stream, CancellationToken, Task<T>> funcWithDirectStreamAsync, CancellationToken cancellationToken);
 
-            protected virtual void OnDisposed()
+            protected virtual void Dispose(bool disposing)
             {
                 // do nothing
             }
@@ -131,8 +123,21 @@ public void Dispose()
 
                 _disposed = true;
 
-                OnDisposed();
+                Dispose(disposing: true);
+                GC.SuppressFinalize(this);
+            }
+
+#if DEBUG
+            ~Connection()
+            {
+                // this can happen if someone kills OOP. 
+                // when that happen, we don't want to crash VS, so this is debug only check
+                if (!Environment.HasShutdownStarted)
+                {
+                    Contract.Requires(false, $@"Should have been disposed!");
+                }
             }
+#endif
         }
     }
 }

src/Workspaces/Core/Portable/Remote/RemoteHostSessionHelpers.cs

@@ -25,7 +25,13 @@ public static async Task<SessionWithSolution> CreateAsync(RemoteHostClient.Conne
 
             try
             {
-                await connection.RegisterPinnedRemotableDataScopeAsync(scope).ConfigureAwait(false);
+                // set connection state for this session.
+                // we might remove this in future. see https://github.com/dotnet/roslyn/issues/24836
+                await connection.InvokeAsync(
+                    WellKnownServiceHubServices.ServiceHubServiceBase_Initialize,
+                    new object[] { scope.SolutionInfo },
+                    cancellationToken).ConfigureAwait(false);
+
                 return sessionWithSolution;
             }
             catch

src/Workspaces/Remote/ServiceHub/Services/CodeAnalysisService_DesignerAttributes.cs

@@ -19,13 +19,13 @@ internal partial class CodeAnalysisService : IRemoteDesignerAttributeService
         /// 
         /// This will be called by ServiceHub/JsonRpc framework
         /// </summary>
-        public Task<DesignerAttributeResult> ScanDesignerAttributesAsync(PinnedSolutionInfo solutionInfo, DocumentId documentId, CancellationToken cancellationToken)
+        public Task<DesignerAttributeResult> ScanDesignerAttributesAsync(DocumentId documentId, CancellationToken cancellationToken)
         {
             return RunServiceAsync(async token =>
             {
                 using (RoslynLogger.LogBlock(FunctionId.CodeAnalysisService_GetDesignerAttributesAsync, documentId.DebugName, token))
                 {
-                    var solution = await GetSolutionAsync(solutionInfo, token).ConfigureAwait(false);
+                    var solution = await GetSolutionAsync(token).ConfigureAwait(false);
                     var document = solution.GetDocument(documentId);
 
                     var service = document.GetLanguageService<IDesignerAttributeService>();

src/Workspaces/Remote/ServiceHub/Services/CodeAnalysisService_TodoComments.cs

@@ -20,14 +20,13 @@ internal partial class CodeAnalysisService : IRemoteTodoCommentService
         /// 
         /// This will be called by ServiceHub/JsonRpc framework
         /// </summary>
-        public async Task<IList<TodoComment>> GetTodoCommentsAsync(
-            PinnedSolutionInfo solutionInfo, DocumentId documentId, IList<TodoCommentDescriptor> tokens, CancellationToken cancellationToken)
+        public async Task<IList<TodoComment>> GetTodoCommentsAsync(DocumentId documentId, IList<TodoCommentDescriptor> tokens, CancellationToken cancellationToken)
         {
             return await RunServiceAsync(async token =>
             {
                 using (RoslynLogger.LogBlock(FunctionId.CodeAnalysisService_GetTodoCommentsAsync, documentId.DebugName, token))
                 {
-                    var solution = await GetSolutionAsync(solutionInfo, token).ConfigureAwait(false);
+                    var solution = await GetSolutionAsync(token).ConfigureAwait(false);
                     var document = solution.GetDocument(documentId);
 
                     var service = document.GetLanguageService<ITodoCommentService>();