Send a Object Explorer session disconnect message on socket exceptions (#739)

* WIP

* WIP 2

* Send disconnect message o binding exception

* Add a try catch around the binding queue error handler

src/Microsoft.SqlTools.ServiceLayer/LanguageServices/BindingQueue.cs

@@ -5,12 +5,15 @@
 
 using System;
 using System.Collections.Generic;
+using System.Data.SqlClient;
+using System.Diagnostics;
+using System.Linq;
+using System.Net.Sockets;
 using System.Threading;
 using System.Threading.Tasks;
-using Microsoft.SqlTools.Utility;
-using System.Linq;
 using Microsoft.SqlTools.ServiceLayer.Utility;
-using System.Diagnostics;
+using Microsoft.SqlTools.Utility;
+
 
 namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
 {    
@@ -33,6 +36,10 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
 
         private Task queueProcessorTask;
 
+        public delegate void UnhandledExceptionDelegate(string connectionKey, Exception ex);
+
+        public event UnhandledExceptionDelegate OnUnhandledException;
+
         /// <summary>
         /// Map from context keys to binding context instances
         /// Internal for testing purposes only
@@ -350,9 +357,26 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
                                     {
                                         Logger.Write(TraceEventType.Error, "Unexpected exception on the binding queue: " + ex.ToString());
                                         if (queueItem.ErrorHandler != null)
+                                        {
+                                            try
                                             {
                                                 result = queueItem.ErrorHandler(ex);
                                             }
+                                            catch (Exception ex2)
+                                            {
+                                                Logger.Write(TraceEventType.Error, "Unexpected exception in binding queue error handler: " + ex2.ToString());
+                                            }
+                                        }
+
+                                        if (IsExceptionOfType(ex, typeof(SqlException)) || IsExceptionOfType(ex, typeof(SocketException)))
+                                        {
+                                            if (this.OnUnhandledException != null)
+                                            {
+                                                this.OnUnhandledException(queueItem.Key, ex);
+                                            }
+
+                                            RemoveBindingContext(queueItem.Key);
+                                        }    
                                     }
                                 });
 
@@ -469,5 +493,10 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
                 }
             }
         }
+        
+        private bool IsExceptionOfType(Exception ex, Type t)
+        {
+            return ex.GetType() == t || (ex.InnerException != null && ex.InnerException.GetType() == t);
+        }
     }
 }

src/Microsoft.SqlTools.ServiceLayer/LanguageServices/ConnectedBindingQueue.cs

@@ -88,9 +88,8 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
         /// Generate a unique key based on the ConnectionInfo object
         /// </summary>
         /// <param name="connInfo"></param>
-        private string GetConnectionContextKey(ConnectionInfo connInfo)
+        internal static string GetConnectionContextKey(ConnectionDetails details)
         {            
-            ConnectionDetails details = connInfo.ConnectionDetails;
             string key = string.Format("{0}_{1}_{2}_{3}",
                 details.ServerName ?? "NULL",
                 details.DatabaseName ?? "NULL",
@@ -108,7 +107,7 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
                 key += "_" + details.GroupId;
             }
 
-            return key;
+            return Uri.EscapeUriString(key);
         }
 
         /// <summary>
@@ -158,7 +157,7 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
 
         public void RemoveBindigContext(ConnectionInfo connInfo)
         {
-            string connectionKey = GetConnectionContextKey(connInfo);
+            string connectionKey = GetConnectionContextKey(connInfo.ConnectionDetails);
             if (BindingContextExists(connectionKey))
             {
                 RemoveBindingContext(connectionKey);
@@ -178,7 +177,7 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
             }
 
             // lookup the current binding context
-            string connectionKey = GetConnectionContextKey(connInfo);
+            string connectionKey = GetConnectionContextKey(connInfo.ConnectionDetails);
             if (BindingContextExists(connectionKey))
             {
                 if (overwrite)

src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/Contracts/CloseSessionRequest.cs

@@ -39,6 +39,29 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Contracts
         public string SessionId { get; set; }
     }
 
+    /// <summary>
+    /// Information returned when a session is disconnected.
+    /// Contains success information and a <see cref="SessionId"/>
+    /// </summary>
+    public class SessionDisconnectedParameters
+    {
+        /// <summary>
+        /// Boolean indicating if the connection was successful
+        /// </summary>
+        public bool Success { get; set; }
+
+        /// <summary>
+        /// Unique ID to use when sending any requests for objects in the
+        /// tree under the node
+        /// </summary>
+        public string SessionId { get; set; }
+
+        /// <summary>
+        /// Error message returned from the engine for a object explorer session failure reason, if any.
+        /// </summary>
+        public string ErrorMessage { get; set; }
+    }
+
     /// <summary>
     /// Establishes an Object Explorer tree session for a specific connection.
     /// This will create a connection to a specific server or database, register
@@ -50,4 +73,14 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Contracts
             RequestType<CloseSessionParams, CloseSessionResponse> Type =
             RequestType<CloseSessionParams, CloseSessionResponse>.Create("objectexplorer/closesession");
     }
+    
+    /// <summary>
+    /// Session disconnected notification
+    /// </summary>
+    public class SessionDisconnectedNotification
+    {
+        public static readonly
+            EventType<SessionDisconnectedParameters> Type =
+            EventType<SessionDisconnectedParameters>.Create("objectexplorer/sessiondisconnected");
+    }
 }

src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/Nodes/TreeNode.cs

@@ -324,9 +324,10 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Nodes
             SmoQueryContext context = this.GetContextAs<SmoQueryContext>();
             bool includeSystemObjects = context != null && context.Database != null ? DatabaseUtils.IsSystemDatabaseConnection(context.Database.Name) : true;
 
-
             if (children.IsPopulating || context == null)
+            {
                 return;
+            }
 
             children.Clear();
             BeginChildrenInit();

src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/ObjectExplorerService.cs

@@ -8,10 +8,12 @@ using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Collections.ObjectModel;
 using System.Composition;
+using System.Diagnostics;
 using System.Globalization;
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.SqlServer.Management.Common;
 using Microsoft.SqlTools.Extensibility;
 using Microsoft.SqlTools.Hosting;
 using Microsoft.SqlTools.Hosting.Protocol;
@@ -25,8 +27,6 @@ using Microsoft.SqlTools.ServiceLayer.SqlContext;
 using Microsoft.SqlTools.ServiceLayer.Utility;
 using Microsoft.SqlTools.ServiceLayer.Workspace;
 using Microsoft.SqlTools.Utility;
-using Microsoft.SqlServer.Management.Common;
-using System.Diagnostics;
 
 namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
 {
@@ -132,6 +132,9 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
         {
             Logger.Write(TraceEventType.Verbose, "ObjectExplorer service initialized");
             this.serviceHost = serviceHost;
+
+            this.ConnectedBindingQueue.OnUnhandledException += OnUnhandledException;
+
             // Register handlers for requests
             serviceHost.SetRequestHandler(CreateSessionRequest.Type, HandleCreateSessionRequest);
             serviceHost.SetRequestHandler(ExpandRequest.Type, HandleExpandRequest);
@@ -513,7 +516,6 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
             }
         }        
 
-
         private async Task<ConnectionCompleteParams> Connect(ConnectParams connectParams, string uri)
         {
             string connectionErrorMessage = string.Empty;
@@ -552,6 +554,18 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
             await serviceHost.SendEvent(CreateSessionCompleteNotification.Type, result);
         }
 
+        internal async Task SendSessionDisconnectedNotification(string uri, bool success, string errorMessage)
+        {
+            Logger.Write(TraceEventType.Information, $"OE session disconnected: {errorMessage}");
+            SessionDisconnectedParameters result = new SessionDisconnectedParameters()
+            {
+                Success = success,
+                ErrorMessage = errorMessage,
+                SessionId = uri
+            };
+            await serviceHost.SendEvent(SessionDisconnectedNotification.Type, result);
+        }
+
         private void RunExpandTask(ObjectExplorerSession session, ExpandParams expandParams, bool forceRefresh = false)
         {
             CancellationTokenSource cancellationTokenSource = new CancellationTokenSource();
@@ -627,22 +641,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
         /// <remarks>Internal for testing purposes only</remarks>
         internal static string GenerateUri(ConnectionDetails details)
         {
-            Validate.IsNotNull("details", details);
+            return ConnectedBindingQueue.GetConnectionContextKey(details);
-            string uri = string.Format(CultureInfo.InvariantCulture, "{0}{1}", uriPrefix, Uri.EscapeUriString(details.ServerName));
-            uri = AppendIfExists(uri, "databaseName", details.DatabaseName);
-            uri = AppendIfExists(uri, "user", details.UserName);
-            uri = AppendIfExists(uri, "groupId", details.GroupId);
-            uri = AppendIfExists(uri, "displayName", details.DatabaseDisplayName);
-            return uri;
-        }
-
-        private static string AppendIfExists(string uri, string propertyName, string propertyValue)
-        {
-            if (!string.IsNullOrEmpty(propertyValue))
-            {
-                uri += string.Format(CultureInfo.InvariantCulture, ";{0}={1}", propertyName, Uri.EscapeUriString(propertyValue));
-            }
-            return uri;
         }      
 
         public IEnumerable<ChildFactory> GetApplicableChildFactories(TreeNode item)
@@ -742,10 +741,37 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
         {
             if (bindingQueue != null)
             {
+                bindingQueue.OnUnhandledException -= OnUnhandledException;
                 bindingQueue.Dispose();
             }            
         }
 
+        private async void OnUnhandledException(string queueKey, Exception ex)
+        {
+            string sessionUri = LookupUriFromQueueKey(queueKey);
+            if (!string.IsNullOrWhiteSpace(sessionUri))
+            {
+                await SendSessionDisconnectedNotification(uri: sessionUri, success: false, errorMessage: ex.ToString());
+            }
+        }
+
+        private string LookupUriFromQueueKey(string queueKey)
+        {
+            foreach (var session in this.sessionMap.Values)
+            {
+                var connInfo = session.ConnectionInfo;
+                if (connInfo != null)
+                {
+                    string currentKey = ConnectedBindingQueue.GetConnectionContextKey(connInfo.ConnectionDetails);
+                    if (queueKey == currentKey)
+                    {
+                        return session.Uri;
+                    }
+                }
+            }
+            return string.Empty;
+        }
+
         internal class ObjectExplorerSession
         {
             private ConnectionService connectionService;