From 87e6da6e4c521e457e038596a16b9921aec51ab6 Mon Sep 17 00:00:00 2001
From: David Fowler <davidfowl@gmail.com>
Date: Mon, 3 Apr 2017 22:21:41 -0700
Subject: [PATCH] Handle exceptions and Cancellation in DisposeAsync (#366)

---
 .../Internal/ConnectionState.cs               | 26 +++++++++--
 .../ConnectionManagerTests.cs                 | 46 +++++++++++++++++++
 2 files changed, 69 insertions(+), 3 deletions(-)

diff --git a/src/Microsoft.AspNetCore.Sockets/Internal/ConnectionState.cs b/src/Microsoft.AspNetCore.Sockets/Internal/ConnectionState.cs
index a8f24bc808..7714d14cff 100644
--- a/src/Microsoft.AspNetCore.Sockets/Internal/ConnectionState.cs
+++ b/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
-            _disposeTcs.TrySetResult(null);
+        private async Task WaitOnTasks(Task applicationTask, Task transportTask)
+        {
+            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
diff --git a/test/Microsoft.AspNetCore.Sockets.Tests/ConnectionManagerTests.cs b/test/Microsoft.AspNetCore.Sockets.Tests/ConnectionManagerTests.cs
index 46549a578c..574f22a152 100644
--- a/test/Microsoft.AspNetCore.Sockets.Tests/ConnectionManagerTests.cs
+++ b/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()
         {