From b45026967ed2c04b667402bb9d85bd6a2246c81c Mon Sep 17 00:00:00 2001
From: Cheena Malhotra <13396919+cheenamalhotra@users.noreply.github.com>
Date: Wed, 31 May 2023 17:15:47 -0700
Subject: [PATCH] Enable connection pooling for languageService (#2080)

---
 .../LanguageServices/LanguageService.cs       |  6 ++---
 .../Hosting/Protocol/Constants.cs             |  3 +++
 .../Connection/ConnectionService.cs           | 27 +++++++++++++------
 .../LanguageServices/LanguageService.cs       |  6 ++---
 4 files changed, 28 insertions(+), 14 deletions(-)

diff --git a/src/Microsoft.Kusto.ServiceLayer/LanguageServices/LanguageService.cs b/src/Microsoft.Kusto.ServiceLayer/LanguageServices/LanguageService.cs
index cf6876ac..d89def83 100644
--- a/src/Microsoft.Kusto.ServiceLayer/LanguageServices/LanguageService.cs
+++ b/src/Microsoft.Kusto.ServiceLayer/LanguageServices/LanguageService.cs
@@ -129,7 +129,7 @@ namespace Microsoft.Kusto.ServiceLayer.LanguageServices
                 if (connectionService == null)
                 {
                     connectionService = ConnectionService.Instance;
-                    connectionService.RegisterConnectedQueue("LanguageService", _bindingQueue);
+                    connectionService.RegisterConnectedQueue(Constants.LanguageServiceFeature, _bindingQueue);
                 }
                 return connectionService;
             }
@@ -557,7 +557,7 @@ namespace Microsoft.Kusto.ServiceLayer.LanguageServices
                         {
                             try
                             {
-                                _bindingQueue.AddConnectionContext(connInfo, true, featureName: "LanguageService", overwrite: true);
+                                _bindingQueue.AddConnectionContext(connInfo, true, featureName: Constants.LanguageServiceFeature, overwrite: true);
                                 RemoveScriptParseInfo(rebuildParams.OwnerUri);
                                 UpdateLanguageServiceOnConnection(connInfo).Wait();
                             }
@@ -677,7 +677,7 @@ namespace Microsoft.Kusto.ServiceLayer.LanguageServices
                 {
                     try
                     {
-                        scriptInfo.ConnectionKey = _bindingQueue.AddConnectionContext(connInfo, true,"languageService");
+                        scriptInfo.ConnectionKey = _bindingQueue.AddConnectionContext(connInfo, true, Constants.LanguageServiceFeature);
                         scriptInfo.IsConnected = _bindingQueue.IsBindingContextConnected(scriptInfo.ConnectionKey);
                     }
                     catch (Exception ex)
diff --git a/src/Microsoft.SqlTools.Hosting/Hosting/Protocol/Constants.cs b/src/Microsoft.SqlTools.Hosting/Hosting/Protocol/Constants.cs
index ff899150..f2b7aabb 100644
--- a/src/Microsoft.SqlTools.Hosting/Hosting/Protocol/Constants.cs
+++ b/src/Microsoft.SqlTools.Hosting/Hosting/Protocol/Constants.cs
@@ -16,6 +16,9 @@ namespace Microsoft.SqlTools.Hosting.Protocol
         public static readonly string DefaultApplicationName = "azdata";
         public static readonly string SqlLoginAuthenticationType = "SqlLogin";
 
+        // Feature names
+        public static readonly string LanguageServiceFeature = "languageService";
+
         static Constants()
         {
             JsonSerializerSettings = new JsonSerializerSettings();
diff --git a/src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionService.cs b/src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionService.cs
index 20bf15ec..d63a5709 100644
--- a/src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionService.cs
+++ b/src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionService.cs
@@ -1329,17 +1329,19 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
         /// <summary>
         /// Build a connection string from a connection details instance
         /// </summary>
-        /// <param name="connectionDetails"></param>
-        public static string BuildConnectionString(ConnectionDetails connectionDetails)
+        /// <param name="connectionDetails">Connection details</param>
+        /// <param name="forceDisablePooling">Whether to disable connection pooling, defaults to true.</param>
+        public static string BuildConnectionString(ConnectionDetails connectionDetails, bool forceDisablePooling = true)
         {
-            return CreateConnectionStringBuilder(connectionDetails).ToString();
+            return CreateConnectionStringBuilder(connectionDetails, forceDisablePooling).ToString();
         }
 
         /// <summary>
         /// Build a connection string builder a connection details instance
         /// </summary>
-        /// <param name="connectionDetails"></param>
-        public static SqlConnectionStringBuilder CreateConnectionStringBuilder(ConnectionDetails connectionDetails)
+        /// <param name="connectionDetails">Connection details</param>
+        /// <param name="forceDisablePooling">Whether to disable connection pooling, defaults to true.</param>
+        public static SqlConnectionStringBuilder CreateConnectionStringBuilder(ConnectionDetails connectionDetails, bool forceDisablePooling = true)
         {
             SqlConnectionStringBuilder connectionBuilder;
 
@@ -1579,7 +1581,10 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
             {
                 connectionBuilder.TypeSystemVersion = connectionDetails.TypeSystemVersion;
             }
-            connectionBuilder.Pooling = false;
+            if (forceDisablePooling)
+            {
+                connectionBuilder.Pooling = false;
+            }
 
             return connectionBuilder;
         }
@@ -1833,17 +1838,23 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
                 bool? originalPersistSecurityInfo = connInfo.ConnectionDetails.PersistSecurityInfo;
                 bool? originalPooling = connInfo.ConnectionDetails.Pooling;
 
+                // allow pooling connections for language service feature to improve intellisense connection retention and performance.
+                bool shouldForceDisablePooling = featureName != Constants.LanguageServiceFeature;
+
                 // increase the connection and command timeout to at least 30 seconds and and build connection string
                 connInfo.ConnectionDetails.ConnectTimeout = Math.Max(30, originalTimeout ?? 0);
                 connInfo.ConnectionDetails.CommandTimeout = Math.Max(30, originalCommandTimeout ?? 0);
                 // enable PersistSecurityInfo to handle issues in SMO where the connection context is lost in reconnections
                 connInfo.ConnectionDetails.PersistSecurityInfo = true;
+
                 // turn off connection pool to avoid hold locks on server resources after calling SqlConnection Close method
-                connInfo.ConnectionDetails.Pooling = false;
+                if (shouldForceDisablePooling) {
+                    connInfo.ConnectionDetails.Pooling = false;
+                }
                 connInfo.ConnectionDetails.ApplicationName = GetApplicationNameWithFeature(connInfo.ConnectionDetails.ApplicationName, featureName);
 
                 // generate connection string
-                string connectionString = ConnectionService.BuildConnectionString(connInfo.ConnectionDetails);
+                string connectionString = ConnectionService.BuildConnectionString(connInfo.ConnectionDetails, shouldForceDisablePooling);
 
                 // restore original values
                 connInfo.ConnectionDetails.ConnectTimeout = originalTimeout;
diff --git a/src/Microsoft.SqlTools.ServiceLayer/LanguageServices/LanguageService.cs b/src/Microsoft.SqlTools.ServiceLayer/LanguageServices/LanguageService.cs
index b85f2680..08793b13 100644
--- a/src/Microsoft.SqlTools.ServiceLayer/LanguageServices/LanguageService.cs
+++ b/src/Microsoft.SqlTools.ServiceLayer/LanguageServices/LanguageService.cs
@@ -172,7 +172,7 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
                 if (connectionService == null)
                 {
                     connectionService = ConnectionService.Instance;
-                    connectionService.RegisterConnectedQueue("LanguageService", bindingQueue);
+                    connectionService.RegisterConnectedQueue(Constants.LanguageServiceFeature, bindingQueue);
                 }
                 return connectionService;
             }
@@ -720,7 +720,7 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
                         {
                             try
                             {
-                                this.BindingQueue.AddConnectionContext(connInfo, featureName: "LanguageService", overwrite: true);
+                                this.BindingQueue.AddConnectionContext(connInfo, featureName: Constants.LanguageServiceFeature, overwrite: true);
                                 RemoveScriptParseInfo(rebuildParams.OwnerUri);
                                 UpdateLanguageServiceOnConnection(connInfo).Wait();
                             }
@@ -1004,7 +1004,7 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
                 {
                     try
                     {
-                        scriptInfo.ConnectionKey = this.BindingQueue.AddConnectionContext(info, "languageService");
+                        scriptInfo.ConnectionKey = this.BindingQueue.AddConnectionContext(info, Constants.LanguageServiceFeature);
                         scriptInfo.IsConnected = this.BindingQueue.IsBindingContextConnected(scriptInfo.ConnectionKey);
                     }
                     catch (Exception ex)