From 283297f455de5d5a22350aba9a512f7db2ffdaf3 Mon Sep 17 00:00:00 2001
From: Mikael Mengistu <mikaelm12@users.noreply.github.com>
Date: Tue, 15 May 2018 23:57:02 -0700
Subject: [PATCH] Remove handlers from HubConnection (#2267)

---
 .../HubConnection.Log.cs                      |  8 +++
 .../HubConnection.cs                          | 10 +++
 .../HubConnectionTests.Protocol.cs            | 66 +++++++++++++++++++
 3 files changed, 84 insertions(+)

diff --git a/src/Microsoft.AspNetCore.SignalR.Client.Core/HubConnection.Log.cs b/src/Microsoft.AspNetCore.SignalR.Client.Core/HubConnection.Log.cs
index 5dbcf2985e..d2791f4d3d 100644
--- a/src/Microsoft.AspNetCore.SignalR.Client.Core/HubConnection.Log.cs
+++ b/src/Microsoft.AspNetCore.SignalR.Client.Core/HubConnection.Log.cs
@@ -171,6 +171,9 @@ namespace Microsoft.AspNetCore.SignalR.Client
             private static readonly Action<ILogger, string, string, Exception> _argumentBindingFailure =
                 LoggerMessage.Define<string, string>(LogLevel.Error, new EventId(57, "ArgumentBindingFailure"), "Failed to bind arguments received in invocation '{InvocationId}' of '{MethodName}'.");
 
+            private static readonly Action<ILogger, string, Exception> _removingHandlers =
+               LoggerMessage.Define<string>(LogLevel.Debug, new EventId(58, "RemovingHandlers"), "Removing handlers for client method '{MethodName}'.");
+
             public static void PreparingNonBlockingInvocation(ILogger logger, string target, int count)
             {
                 _preparingNonBlockingInvocation(logger, target, count, null);
@@ -360,6 +363,11 @@ namespace Microsoft.AspNetCore.SignalR.Client
                 _registeringHandler(logger, methodName, null);
             }
 
+            public static void RemovingHandlers(ILogger logger, string methodName)
+            {
+                _removingHandlers(logger, methodName, null);
+            }
+
             public static void Starting(ILogger logger)
             {
                 _starting(logger, null);
diff --git a/src/Microsoft.AspNetCore.SignalR.Client.Core/HubConnection.cs b/src/Microsoft.AspNetCore.SignalR.Client.Core/HubConnection.cs
index f9d02156aa..568bbc35c4 100644
--- a/src/Microsoft.AspNetCore.SignalR.Client.Core/HubConnection.cs
+++ b/src/Microsoft.AspNetCore.SignalR.Client.Core/HubConnection.cs
@@ -175,6 +175,16 @@ namespace Microsoft.AspNetCore.SignalR.Client
 
             return new Subscription(invocationHandler, invocationList);
         }
+        /// <summary>
+        /// Removes all handlers associated with the method with the specified method name.
+        /// </summary>
+        /// <param name="methodName">The name of the hub method from which handlers are being removed</param>
+        public void Remove(string methodName)
+        {
+            CheckDisposed();
+            Log.RemovingHandlers(_logger, methodName);
+            _handlers.TryRemove(methodName, out _);
+        }
 
         /// <summary>
         /// Invokes a streaming hub method on the server using the specified method name, return type and arguments.
diff --git a/test/Microsoft.AspNetCore.SignalR.Client.Tests/HubConnectionTests.Protocol.cs b/test/Microsoft.AspNetCore.SignalR.Client.Tests/HubConnectionTests.Protocol.cs
index 3bd1dd833f..99ec2eba60 100644
--- a/test/Microsoft.AspNetCore.SignalR.Client.Tests/HubConnectionTests.Protocol.cs
+++ b/test/Microsoft.AspNetCore.SignalR.Client.Tests/HubConnectionTests.Protocol.cs
@@ -379,6 +379,72 @@ namespace Microsoft.AspNetCore.SignalR.Client.Tests
                 }
             }
 
+            [Fact]
+            public async Task HandlerIsRemovedProperlyWithOff()
+            {
+                var connection = new TestConnection();
+                var hubConnection = CreateHubConnection(connection);
+                var handlerCalled = new TaskCompletionSource<int>();
+                try
+                {
+                    await hubConnection.StartAsync().OrTimeout();
+
+                    hubConnection.On<int>("Foo", (val) =>
+                    {
+                        handlerCalled.TrySetResult(val);
+                    });
+
+                    hubConnection.Remove("Foo");
+                    await connection.ReceiveJsonMessage(new { invocationId = "1", type = 1, target = "Foo", arguments = 1 }).OrTimeout();
+                    var handlerTask = handlerCalled.Task;
+
+                    // We expect the handler task to timeout since the handler has been removed with the call to Remove("Foo")
+                    var ex = Assert.ThrowsAsync<TimeoutException>(async () => await handlerTask.OrTimeout(2000));
+
+                    // Ensure that the task from the WhenAny is not the handler task
+                    Assert.False(handlerCalled.Task.IsCompleted);
+                }
+                finally
+                {
+                    await hubConnection.DisposeAsync().OrTimeout();
+                    await connection.DisposeAsync().OrTimeout();
+                }
+            }
+
+            [Fact]
+            public async Task DisposingSubscriptionAfterCallingRemoveHandlerDoesntFail()
+            {
+                var connection = new TestConnection();
+                var hubConnection = CreateHubConnection(connection);
+                var handlerCalled = new TaskCompletionSource<int>();
+                try
+                {
+                    await hubConnection.StartAsync().OrTimeout();
+
+                    var subscription = hubConnection.On<int>("Foo", (val) =>
+                    {
+                        handlerCalled.TrySetResult(val);
+                    });
+
+                    hubConnection.Remove("Foo");
+                    await connection.ReceiveJsonMessage(new { invocationId = "1", type = 1, target = "Foo", arguments = 1 }).OrTimeout();
+                    var handlerTask = handlerCalled.Task;
+
+                    subscription.Dispose();
+
+                    // We expect the handler task to timeout since the handler has been removed with the call to Remove("Foo")
+                    var ex = Assert.ThrowsAsync<TimeoutException>(async () => await handlerTask.OrTimeout(2000));
+
+                    // Ensure that the task from the WhenAny is not the handler task
+                    Assert.False(handlerCalled.Task.IsCompleted);
+                }
+                finally
+                {
+                    await hubConnection.DisposeAsync().OrTimeout();
+                    await connection.DisposeAsync().OrTimeout();
+                }
+            }
+
             [Fact]
             public async Task AcceptsPingMessages()
             {