Query Execution: Better exception handling on unawaited async tasks (#502)

* WIP

* This code makes it work!

* Adding similar exception hanling behavior to saving result sets

* Adding unit tests for new extension methods
Auto-cleanup of proj file whitespace

* Implementing changes as per code review comments

src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Query.cs

@@ -1,6 +1,7 @@
 // 
 // Copyright (c) Microsoft. All rights reserved.
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
+//
 
 using System;
 using System.Data.Common;
@@ -17,6 +18,7 @@ using Microsoft.SqlTools.ServiceLayer.SqlContext;
 using Microsoft.SqlTools.Utility;
 using Microsoft.SqlTools.ServiceLayer.BatchParser.ExecutionEngineCode;
 using System.Collections.Generic;
+using Microsoft.SqlTools.ServiceLayer.Utility;
 
 namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
 {
@@ -25,11 +27,35 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
     /// </summary>
     public class Query : IDisposable
     {
+        #region Constants
+        
         /// <summary>
         /// "Error" code produced by SQL Server when the database context (name) for a connection changes.
         /// </summary>
         private const int DatabaseContextChangeErrorNumber = 5701;
         
+        /// <summary>
+        /// ON keyword
+        /// </summary>
+        private const string On = "ON";
+
+        /// <summary>
+        /// OFF keyword
+        /// </summary>
+        private const string Off = "OFF";
+        
+        /// <summary>
+        /// showplan_xml statement
+        /// </summary>
+        private const string SetShowPlanXml = "SET SHOWPLAN_XML {0}";
+
+        /// <summary>
+        /// statistics xml statement
+        /// </summary>
+        private const string SetStatisticsXml = "SET STATISTICS XML {0}";
+        
+        #endregion
+
         #region Member Variables
 
         /// <summary>
@@ -58,26 +84,6 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         /// </summary>
         private string newDatabaseName;        
 
-        /// <summary>
-        /// ON keyword
-        /// </summary>
-        private const string On = "ON";
-
-        /// <summary>
-        /// OFF keyword
-        /// </summary>
-        private const string Off = "OFF";
-
-        /// <summary>
-        /// showplan_xml statement
-        /// </summary>
-        private const string SetShowPlanXml = "SET SHOWPLAN_XML {0}";
-
-        /// <summary>
-        /// statistics xml statement
-        /// </summary>
-        private const string SetStatisticsXml = "SET STATISTICS XML {0}";
-
         #endregion
 
         /// <summary>
@@ -139,6 +145,19 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
 
         #region Events
 
+        /// <summary>
+        /// Delegate type for callback when a query completes or fails
+        /// </summary>
+        /// <param name="query">The query that completed</param>
+        public delegate Task QueryAsyncEventHandler(Query query);
+        
+        /// <summary>
+        /// Delegate type for callback when a query fails
+        /// </summary>
+        /// <param name="query">Query that raised the event</param>
+        /// <param name="exception">Exception that caused the query to fail</param>
+        public delegate Task QueryAsyncErrorEventHandler(Query query, Exception exception);
+        
         /// <summary>
         /// Event to be called when a batch is completed.
         /// </summary>
@@ -154,12 +173,6 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         /// </summary>
         public event Batch.BatchAsyncEventHandler BatchStarted;
 
-        /// <summary>
-        /// Delegate type for callback when a query connection fails
-        /// </summary>
-        /// <param name="message">Error message for the failing query</param>
-        public delegate Task QueryAsyncErrorEventHandler(Query q, Exception e);
-
         /// <summary>
         /// Callback for when the query has completed successfully
         /// </summary>
@@ -179,26 +192,20 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
 
         #region Properties
 
-        /// <summary>
-        /// Delegate type for callback when a query completes or fails
-        /// </summary>
-        /// <param name="q">The query that completed</param>
-        public delegate Task QueryAsyncEventHandler(Query q);
-
         /// <summary>
         /// The batches which should run before the user batches 
         /// </summary>
-        internal List<Batch> BeforeBatches { get; set; }
+        private List<Batch> BeforeBatches { get; }
 
         /// <summary>
         /// The batches underneath this query
         /// </summary>
-        internal Batch[] Batches { get; set; }
+        internal Batch[] Batches { get; }
 
         /// <summary>
         /// The batches which should run after the user batches 
         /// </summary>
-        internal List<Batch> AfterBatches { get; set; }
+        internal List<Batch> AfterBatches { get; }
 
         /// <summary>
         /// The summaries of the batches underneath this query
@@ -243,7 +250,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         /// <summary>
         /// The text of the query to execute
         /// </summary>
-        public string QueryText { get; set; }
+        public string QueryText { get; }
 
         #endregion
 
@@ -269,7 +276,11 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         /// </summary>
         public void Execute()
         {
-            ExecutionTask = Task.Run(ExecuteInternal);
+            ExecutionTask = Task.Run(ExecuteInternal)
+                .ContinueWithOnFaulted(t =>
+                {
+                    QueryFailed?.Invoke(this, t.Exception).Wait();
+                });
         }
 
         /// <summary>
@@ -338,34 +349,36 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         /// </summary>
         private async Task ExecuteInternal()
         {
-            // Mark that we've internally executed
+            ReliableSqlConnection sqlConn = null;
-            hasExecuteBeenCalled = true;
-
-            // Don't actually execute if there aren't any batches to execute
-            if (Batches.Length == 0)
-            {
-                if (BatchMessageSent != null)
-                {
-                    await BatchMessageSent(new ResultMessage(SR.QueryServiceCompletedSuccessfully, false, null));
-                }
-                if (QueryCompleted != null)
-                {
-                    await QueryCompleted(this);
-                }
-                return;
-            }
-
-            // Locate and setup the connection
-            DbConnection queryConnection = await ConnectionService.Instance.GetOrOpenConnection(editorConnection.OwnerUri, ConnectionType.Query);
-            ReliableSqlConnection sqlConn = queryConnection as ReliableSqlConnection;
-            if (sqlConn != null)
-            {
-                // Subscribe to database informational messages
-                sqlConn.GetUnderlyingConnection().InfoMessage += OnInfoMessage;
-            }
-
             try
             {
+                // Mark that we've internally executed
+                hasExecuteBeenCalled = true;
+    
+                // Don't actually execute if there aren't any batches to execute
+                if (Batches.Length == 0)
+                {
+                    if (BatchMessageSent != null)
+                    {
+                        await BatchMessageSent(new ResultMessage(SR.QueryServiceCompletedSuccessfully, false, null));
+                    }
+                    if (QueryCompleted != null)
+                    {
+                        await QueryCompleted(this);
+                    }
+                    return;
+                }
+    
+                // Locate and setup the connection
+                DbConnection queryConnection = await ConnectionService.Instance.GetOrOpenConnection(editorConnection.OwnerUri, ConnectionType.Query);
+                sqlConn = queryConnection as ReliableSqlConnection;
+                if (sqlConn != null)
+                {
+                    // Subscribe to database informational messages
+                    sqlConn.GetUnderlyingConnection().InfoMessage += OnInfoMessage;
+                }
+
+            
                 // Execute beforeBatches synchronously, before the user defined batches 
                 foreach (Batch b in BeforeBatches)
                 {
@@ -405,16 +418,18 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             }
             finally
             {
+                // Remove the message handler from the connection
                 if (sqlConn != null)
                 {
                     // Subscribe to database informational messages
                     sqlConn.GetUnderlyingConnection().InfoMessage -= OnInfoMessage;
                 }
-            }
                 
-            if (newDatabaseName != null)
+                // If any message notified us we had changed databases, then we must let the connection service know 
-            {
+                if (newDatabaseName != null)
-                ConnectionService.Instance.ChangeConnectionDatabaseContext(editorConnection.OwnerUri, newDatabaseName);
+                {
+                    ConnectionService.Instance.ChangeConnectionDatabaseContext(editorConnection.OwnerUri, newDatabaseName);
+                }
             }
         }
 

src/Microsoft.SqlTools.ServiceLayer/QueryExecution/ResultSet.cs

@@ -505,8 +505,14 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                 }
             });
             
+            // Add exception handling to the save task
+            Task taskWithHandling = saveAsTask.ContinueWithOnFaulted(t =>
+            {
+                failureHandler?.Invoke(saveParams, t.Exception.Message).Wait();
+            });
+
             // If saving the task fails, return a failure
-            if (!SaveTasks.TryAdd(saveParams.FilePath, saveAsTask))
+            if (!SaveTasks.TryAdd(saveParams.FilePath, taskWithHandling))
             {
                 throw new InvalidOperationException(SR.QueryServiceSaveAsMiscStartingError);
             }
@@ -537,7 +543,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             if (!SaveTasks.IsEmpty)
             {
                 // Wait for tasks to finish before disposing ResultSet
-                Task.WhenAll(SaveTasks.Values.ToArray()).ContinueWith((antecedent) =>
+                Task.WhenAll(SaveTasks.Values.ToArray()).ContinueWith(antecedent =>
                 {
                     if (disposing)
                     {

src/Microsoft.SqlTools.ServiceLayer/Utility/TaskExtensions.cs

@@ -0,0 +1,50 @@
+// 
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+//
+
+using System;
+using System.Text;
+using System.Threading.Tasks;
+using Microsoft.SqlTools.Utility;
+
+namespace Microsoft.SqlTools.ServiceLayer.Utility
+{
+    public static class TaskExtensions
+    {
+        /// <summary>
+        /// Adds handling to check the Exception field of a task and log it if the task faulted
+        /// </summary>
+        /// <remarks>
+        /// This will effectively swallow exceptions in the task chain. 
+        /// </remarks>
+        /// <param name="task">The task to continue</param>
+        /// <param name="continuationAction">
+        /// An optional operation to perform after exception handling has occurred
+        /// </param>
+        /// <returns>Task with exception handling on continuation</returns>
+        public static Task ContinueWithOnFaulted(this Task task, Action<Task> continuationAction)
+        {
+            return task.ContinueWith(t =>
+            {
+                // If the task hasn't faulted or has an exception, skip processing
+                if (!t.IsFaulted || t.Exception == null)
+                {
+                    return;
+                }
+                
+                // Construct an error message for an aggregate exception and log it
+                StringBuilder sb = new StringBuilder("Unhandled exception(s) in async task:");
+                foreach (Exception e in task.Exception.InnerExceptions)
+                {
+                    sb.AppendLine($"{e.GetType().Name}: {e.Message}");
+                    sb.AppendLine(e.StackTrace);
+                }
+                Logger.Write(LogLevel.Error, sb.ToString());
+
+                // Run the continuation task that was provided
+                continuationAction?.Invoke(t);
+            });
+        }
+    }
+}

test/Microsoft.SqlTools.ServiceLayer.UnitTests/Microsoft.SqlTools.ServiceLayer.UnitTests.csproj

@@ -2,8 +2,8 @@
   <PropertyGroup Label="Configuration">
     <OutputType>Exe</OutputType>
     <TargetFramework>netcoreapp2.0</TargetFramework>
-	<GenerateAssemblyInfo>false</GenerateAssemblyInfo>
+    <GenerateAssemblyInfo>false</GenerateAssemblyInfo>
-	<DefineConstants>$(DefineConstants);NETCOREAPP1_0</DefineConstants>
+    <DefineConstants>$(DefineConstants);NETCOREAPP1_0</DefineConstants>
     <IsPackable>false</IsPackable>
     <ApplicationIcon />
     <StartupObject />
@@ -12,8 +12,8 @@
     <PackageReference Include="Microsoft.NET.Test.Sdk" Version="15.3.0" />
     <PackageReference Include="xunit" Version="2.2.0" />
     <PackageReference Include="xunit.runner.visualstudio" Version="2.2.0" />
-	  <PackageReference Include="System.Data.SqlClient" Version="4.4.0" />
+    <PackageReference Include="System.Data.SqlClient" Version="4.4.0" />
-  	<PackageReference Include="Microsoft.SqlServer.Smo" Version="140.2.6" />
+    <PackageReference Include="Microsoft.SqlServer.Smo" Version="140.2.6" />
   </ItemGroup>
   <ItemGroup>
     <Reference Include="Newtonsoft.Json">

test/Microsoft.SqlTools.ServiceLayer.UnitTests/Utility/TaskExtensionTests.cs

@@ -0,0 +1,53 @@
+//
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+//
+
+using System;
+using System.Threading.Tasks;
+using Microsoft.SqlTools.ServiceLayer.Utility;
+using Xunit;
+
+namespace Microsoft.SqlTools.ServiceLayer.UnitTests.Utility
+{
+    public class TaskExtensionTests
+    {
+        [Fact]
+        public async Task ContinueWithOnFaultedNullContinuation()
+        {
+            // Setup: Create a task that will definitely fault
+            Task failureTask = new Task(() => throw new Exception("It fail!"));
+            
+            // If: I continue on fault and start the task
+            Task continuationTask = failureTask.ContinueWithOnFaulted(null);
+            failureTask.Start();
+            await continuationTask;
+
+            // Then: The task should have completed without fault
+            Assert.Equal(TaskStatus.RanToCompletion, continuationTask.Status);
+        }
+
+        [Fact]
+        public async Task ContinueWithOnFaultedContinuatation()
+        {
+            // Setup: 
+            // ... Create a new task that will definitely fault
+            Task failureTask = new Task(() => throw new Exception("It fail!"));
+            
+            // ... Create a quick continuation task that will signify if it's been called
+            Task providedTask = null;
+
+            // If: I continue on fault, with a continuation task
+            Task continuationTask = failureTask.ContinueWithOnFaulted(task => { providedTask = task; });
+            failureTask.Start();
+            await continuationTask;
+            
+            // Then:
+            // ... The task should have completed without fault
+            Assert.Equal(TaskStatus.RanToCompletion, continuationTask.Status);
+            
+            // ... The continuation action should have been called with the original failure task
+            Assert.Equal(failureTask, providedTask);
+        }
+    }
+}