Handle exceptions and Cancellation in DisposeAsync (#366)

src/Microsoft.AspNetCore.Sockets/Internal/ConnectionState.cs

@@ -72,7 +72,7 @@ namespace Microsoft.AspNetCore.Sockets.Internal
                     var applicationTask = ApplicationTask ?? TaskCache.CompletedTask;
                     var transportTask = TransportTask ?? TaskCache.CompletedTask;
 
-                    disposeTask = Task.WhenAll(applicationTask, transportTask);
+                    disposeTask = WaitOnTasks(applicationTask, transportTask);
                 }
             }
             finally
@@ -82,9 +82,29 @@ namespace Microsoft.AspNetCore.Sockets.Internal
 
             // REVIEW: Add a timeout so we don't wait forever
             await disposeTask;
+        }
 
-            // Notify all waiters that we're done disposing
+        private async Task WaitOnTasks(Task applicationTask, Task transportTask)
-            _disposeTcs.TrySetResult(null);
+        {
+            try
+            {
+                await Task.WhenAll(applicationTask, transportTask);
+
+                // Notify all waiters that we're done disposing
+                _disposeTcs.TrySetResult(null);
+            }
+            catch (OperationCanceledException)
+            {
+                _disposeTcs.TrySetCanceled();
+
+                throw;
+            }
+            catch (Exception ex)
+            {
+                _disposeTcs.TrySetException(ex);
+
+                throw;
+            }
         }
 
         public enum ConnectionStatus

test/Microsoft.AspNetCore.Sockets.Tests/ConnectionManagerTests.cs

@@ -1,6 +1,7 @@
 // Copyright (c) .NET Foundation. 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.Threading.Tasks;
 using Microsoft.AspNetCore.SignalR.Tests.Common;
 using Microsoft.AspNetCore.Sockets.Internal;
@@ -121,6 +122,51 @@ namespace Microsoft.AspNetCore.Sockets.Tests
             await Task.WhenAll(firstTask, secondTask).OrTimeout();
         }
 
+        [Fact]
+        public async Task DisposingConnectionMultipleGetsExceptionFromTransportOrApp()
+        {
+            var connectionManager = CreateConnectionManager();
+            var state = connectionManager.CreateConnection();
+            var tcs = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
+
+            state.ApplicationTask = tcs.Task;
+            state.TransportTask = tcs.Task;
+
+            var firstTask = state.DisposeAsync();
+            var secondTask = state.DisposeAsync();
+            Assert.False(firstTask.IsCompleted);
+            Assert.False(secondTask.IsCompleted);
+
+            tcs.TrySetException(new InvalidOperationException("Error"));
+
+            var exception = await Assert.ThrowsAsync<InvalidOperationException>(async () => await firstTask.OrTimeout());
+            Assert.Equal("Error", exception.Message);
+
+            exception = await Assert.ThrowsAsync<InvalidOperationException>(async () => await secondTask.OrTimeout());
+            Assert.Equal("Error", exception.Message);
+        }
+
+        [Fact]
+        public async Task DisposingConnectionMultipleGetsCancellation()
+        {
+            var connectionManager = CreateConnectionManager();
+            var state = connectionManager.CreateConnection();
+            var tcs = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
+
+            state.ApplicationTask = tcs.Task;
+            state.TransportTask = tcs.Task;
+
+            var firstTask = state.DisposeAsync();
+            var secondTask = state.DisposeAsync();
+            Assert.False(firstTask.IsCompleted);
+            Assert.False(secondTask.IsCompleted);
+
+            tcs.TrySetCanceled();
+
+            await Assert.ThrowsAsync<TaskCanceledException>(async () => await firstTask.OrTimeout());
+            await Assert.ThrowsAsync<TaskCanceledException>(async () => await secondTask.OrTimeout());
+        }
+
         [Fact]
         public async Task DisposeInactiveConnection()
         {