Fix serialization tests & don't block thread (#846)

* Fix serialization tests & don't block thread
- Fixed potential null ref when closing streams
- Always clean up serialization queue if an error occurs
- Stop blocking dispatcher thread by not awaiting task that processes the message
- Improved error logging in EventFlowValidator to help debug issues
- Close stream on exception

src/Microsoft.SqlTools.Hosting/Hosting/Contracts/Error.cs

@@ -20,5 +20,10 @@ namespace Microsoft.SqlTools.Hosting.Contracts
         /// Error message
         /// </summary>
         public string Message { get; set; }
+
+        public override string ToString()
+        {
+            return $"Error(Code={Code},Message='{Message}')";
+        }
     }
 }

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

@@ -14,6 +14,7 @@ using Microsoft.SqlTools.Hosting;
 using Microsoft.SqlTools.Hosting.Protocol;
 using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts;
 using Microsoft.SqlTools.ServiceLayer.QueryExecution.DataStorage;
+using Microsoft.SqlTools.ServiceLayer.Utility;
 using Microsoft.SqlTools.Utility;
 
 
@@ -40,93 +41,111 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         /// <summary>
         /// Begin to process request to save a resultSet to a file in CSV format
         /// </summary>
-        internal async Task HandleSerializeStartRequest(SerializeDataStartRequestParams serializeParams,
+        internal Task HandleSerializeStartRequest(SerializeDataStartRequestParams serializeParams,
             RequestContext<SerializeDataResult> requestContext)
+        {
+            // Run in separate thread so that message thread isn't held up by a potentially time consuming file write
+            Task.Run(async () => {
+                await RunSerializeStartRequest(serializeParams, requestContext);
+            }).ContinueWithOnFaulted(async t => await SendErrorAndCleanup(serializeParams?.FilePath, requestContext, t.Exception));
+            return Task.CompletedTask;
+        }
+
+        internal async Task RunSerializeStartRequest(SerializeDataStartRequestParams serializeParams, RequestContext<SerializeDataResult> requestContext)
         {
             try
             {
+                // Verify we have sensible inputs and there isn't a task running for this file already
                 Validate.IsNotNull(nameof(serializeParams), serializeParams);
                 Validate.IsNotNullOrWhitespaceString("FilePath", serializeParams.FilePath);
 
                 DataSerializer serializer = null;
-                bool hasSerializer = inProgressSerializations.TryGetValue(serializeParams.FilePath, out serializer);
+                if (inProgressSerializations.TryGetValue(serializeParams.FilePath, out serializer))
-                if (hasSerializer)
                 {
+                    // Cannot proceed as there is an in progress serialization happening
                     throw new Exception(SR.SerializationServiceRequestInProgress(serializeParams.FilePath));
                 }
 
+                // Create a new serializer, save for future calls if needed, and write the request out
                 serializer = new DataSerializer(serializeParams);
                 if (!serializeParams.IsLastBatch)
                 {
                     inProgressSerializations.AddOrUpdate(serializer.FilePath, serializer, (key, old) => serializer);
                 }
-                Func<Task<SerializeDataResult>> writeData = () =>
+                
-                {
+                Logger.Write(TraceEventType.Verbose, "HandleSerializeStartRequest");
-                    return Task.Factory.StartNew(() =>
+                SerializeDataResult result = serializer.ProcessRequest(serializeParams);
-                    {
+                await requestContext.SendResult(result);
-                        var result = serializer.ProcessRequest(serializeParams);
-                        return result;
-                    });
-                };
-                await HandleRequest(writeData, requestContext, "HandleSerializeStartRequest");
             }
             catch (Exception ex)
             {
-                await requestContext.SendError(ex.Message);
+                await SendErrorAndCleanup(serializeParams.FilePath, requestContext, ex);
             }
         }
 
+        private async Task SendErrorAndCleanup(string filePath, RequestContext<SerializeDataResult> requestContext, Exception ex)
+        {
+            if (filePath != null)
+            {
+                try
+                {
+                    DataSerializer removed;
+                    inProgressSerializations.TryRemove(filePath, out removed);
+                    if (removed != null)
+                    {
+                        // Flush any contents to disk and remove the writer
+                        removed.CloseStreams();
+                    }
+                }
+                catch
+                {
+                    // Do not care if there was an error removing this, must always delete if something failed
+                }
+            }
+            await requestContext.SendError(ex.Message);
+        }
+
         /// <summary>
         /// Process request to save a resultSet to a file in CSV format
         /// </summary>
-        internal async Task HandleSerializeContinueRequest(SerializeDataContinueRequestParams serializeParams,
+        internal Task HandleSerializeContinueRequest(SerializeDataContinueRequestParams serializeParams,
             RequestContext<SerializeDataResult> requestContext)
+        {
+            // Run in separate thread so that message thread isn't held up by a potentially time consuming file write
+            Task.Run(async () =>
+            {
+                await RunSerializeContinueRequest(serializeParams, requestContext);
+            }).ContinueWithOnFaulted(async t => await SendErrorAndCleanup(serializeParams?.FilePath, requestContext, t.Exception));
+            return Task.CompletedTask;
+        }
+
+        internal async Task RunSerializeContinueRequest(SerializeDataContinueRequestParams serializeParams, RequestContext<SerializeDataResult> requestContext)
         {
             try
             {
+                // Verify we have sensible inputs and some data has already been sent for the file
                 Validate.IsNotNull(nameof(serializeParams), serializeParams);
                 Validate.IsNotNullOrWhitespaceString("FilePath", serializeParams.FilePath);
 
                 DataSerializer serializer = null;
-                bool hasSerializer = inProgressSerializations.TryGetValue(serializeParams.FilePath, out serializer);
+                if (!inProgressSerializations.TryGetValue(serializeParams.FilePath, out serializer))
-                if (!hasSerializer)
                 {
                     throw new Exception(SR.SerializationServiceRequestNotFound(serializeParams.FilePath));
                 }
 
-                Func<Task<SerializeDataResult>> writeData = () =>
+                // Write to file and cleanup if needed
-                {
+                Logger.Write(TraceEventType.Verbose, "HandleSerializeContinueRequest");
-                    return Task.Factory.StartNew(() =>
+                SerializeDataResult result = serializer.ProcessRequest(serializeParams);
-                    {
-                        var result = serializer.ProcessRequest(serializeParams);
                 if (serializeParams.IsLastBatch)
                 {
                     // Cleanup the serializer
                     this.inProgressSerializations.TryRemove(serializer.FilePath, out serializer);
                 }
-                        return result;
-                    });
-                };
-                await HandleRequest(writeData, requestContext, "HandleSerializeContinueRequest");
-            }
-            catch (Exception ex)
-            {
-                await requestContext.SendError(ex.Message);
-            }
-        }
-
-        private async Task HandleRequest<T>(Func<Task<T>> handler, RequestContext<T> requestContext, string requestType)
-        {
-            Logger.Write(TraceEventType.Verbose, requestType);
-
-            try
-            {
-                T result = await handler();
                 await requestContext.SendResult(result);
             }
             catch (Exception ex)
             {
-                await requestContext.SendError(ex.Message);
+                await SendErrorAndCleanup(serializeParams.FilePath, requestContext, ex);
             }
         }
     }
@@ -242,9 +261,13 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                 this.writer = factory.GetWriter(requestParams.FilePath);
             }
         }
-        private void CloseStreams()
+        public void CloseStreams()
+        {
+            if (this.writer != null)
             {
                 this.writer.Dispose();
+                this.writer = null;
+            }
         }
 
         private SaveResultsAsJsonRequestParams CreateJsonRequestParams()

test/Microsoft.SqlTools.ServiceLayer.Test.Common/RequestContextMocking/EventFlowValidator.cs

@@ -153,7 +153,8 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.Common.RequestContextMocking
                 ReceivedEvent received = ReceivedEvents[i];
 
                 // Step 1) Make sure the event type matches
-                Assert.Equal(expected.EventType, received.EventType);
+                Assert.True(expected.EventType.Equals(received.EventType),
+                    string.Format("Expected EventType {0} but got {1}. Received object is {2}", expected.EventType, received.EventType, received.EventObject.ToString()));
                 
                 // Step 2) Make sure the param type matches
                 Assert.True( expected.ParamType == received.EventObject.GetType()

test/Microsoft.SqlTools.ServiceLayer.UnitTests/QueryExecution/SaveResults/SerializationServiceTests.cs

@@ -92,7 +92,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.QueryExecution.SaveResults
                     .AddStandardResultValidator()
                     .Complete();
 
-                await SerializationService.HandleSerializeStartRequest(saveParams, efv.Object);
+                await SerializationService.RunSerializeStartRequest(saveParams, efv.Object);
 
                 // Then:
                 // ... There should not have been an error
@@ -189,7 +189,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.QueryExecution.SaveResults
                 .AddStandardResultValidator()
                 .Complete();
 
-            await SerializationService.HandleSerializeStartRequest(request1, efv.Object);
+            await SerializationService.RunSerializeStartRequest(request1, efv.Object);
         
             // Then:
             // ... There should not have been an error
@@ -202,7 +202,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.QueryExecution.SaveResults
                 .AddStandardResultValidator()
                 .Complete();
 
-            await SerializationService.HandleSerializeContinueRequest(request1, efv.Object);
+            await SerializationService.RunSerializeContinueRequest(request1, efv.Object);
 
             // Then:
             // ... There should not have been an error
@@ -260,10 +260,10 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.QueryExecution.SaveResults
         private static void AssertLineEquals(string line, string[] expected)
         {
             var actual = line.Split(',');
-            Assert.True(actual.Length == expected.Length, string.Format("Line '{0}' does not match values {1}", line, string.Join(",", expected)));
+            Assert.True(actual.Length == expected.Length, $"Line '{line}' does not match values {string.Join(",", expected)}");
             for (int i = 0; i < actual.Length; i++)
             {
-                Assert.True(expected[i] == actual[i], string.Format("Line '{0}' does not match values '{1}' as '{2}' does not equal '{3}'", line, string.Join(",", expected), expected[i], actual[i]));
+                Assert.True(expected[i] == actual[i], $"Line '{line}' does not match values '{string.Join(",", expected)}' as '{expected[i]}' does not equal '{actual[i]}'");
             }
         }