Adding the Subscription type to remove specific handlers (#2646)

clients/java/signalr/src/main/java/HubConnection.java

@@ -88,43 +88,49 @@ public class HubConnection {
         transport.send(message);
     }
 
-    public void on(String target, Action callback) {
+    public Subscription on(String target, Action callback) {
-        ActionBase ac = args -> callback.invoke();
+        ActionBase action = args -> callback.invoke();
-        handlers.put(target, ac);
+        handlers.put(target, action);
+        return new Subscription(handlers, action, target);
     }
 
-    public <T1> void on(String target, Action1<T1> callback, Class<T1> param1) {
+    public <T1> Subscription on(String target, Action1<T1> callback, Class<T1> param1) {
-        ActionBase ac = params -> callback.invoke(param1.cast(params[0]));
+        ActionBase action = params -> callback.invoke(param1.cast(params[0]));
-        handlers.put(target, ac);
+        handlers.put(target, action);
+        return new Subscription(handlers, action, target);
     }
 
-    public <T1, T2> void on(String target, Action2<T1,T2> callback, Class<T1> param1, Class<T2> param2) {
+    public <T1, T2> Subscription on(String target, Action2<T1, T2> callback, Class<T1> param1, Class<T2> param2) {
         ActionBase action = params -> {
             callback.invoke(param1.cast(params[0]), param2.cast(params[1]));
         };
         handlers.put(target, action);
+        return new Subscription(handlers, action, target);
     }
 
-    public <T1, T2, T3> void on(String target, Action3<T1,T2, T3> callback, Class<T1> param1, Class<T2> param2, Class<T3> param3) {
+    public <T1, T2, T3> Subscription on(String target, Action3<T1, T2, T3> callback, Class<T1> param1, Class<T2> param2, Class<T3> param3) {
         ActionBase action = params -> {
             callback.invoke(param1.cast(params[0]), param2.cast(params[1]), param3.cast(params[2]));
         };
         handlers.put(target, action);
+        return new Subscription(handlers, action, target);
     }
 
-    public <T1, T2, T3, T4> void on(String target, Action4<T1,T2, T3, T4> callback, Class<T1> param1, Class<T2> param2, Class<T3> param3, Class<T4> param4) {
+    public <T1, T2, T3, T4> Subscription on(String target, Action4<T1, T2, T3, T4> callback, Class<T1> param1, Class<T2> param2, Class<T3> param3, Class<T4> param4) {
         ActionBase action = params -> {
             callback.invoke(param1.cast(params[0]), param2.cast(params[1]), param3.cast(params[2]), param4.cast(params[3]));
         };
         handlers.put(target, action);
+        return new Subscription(handlers, action, target);
     }
 
-    public <T1, T2, T3, T4,T5> void on(String target, Action5<T1,T2, T3, T4, T5> callback, Class<T1> param1, Class<T2> param2, Class<T3> param3, Class<T4> param4, Class<T5> param5) {
+    public <T1, T2, T3, T4, T5> Subscription on(String target, Action5<T1, T2, T3, T4, T5> callback, Class<T1> param1, Class<T2> param2, Class<T3> param3, Class<T4> param4, Class<T5> param5) {
         ActionBase action = params -> {
             callback.invoke(param1.cast(params[0]), param2.cast(params[1]), param3.cast(params[2]), param4.cast(params[3]),
                     param5.cast(params[4]));
         };
         handlers.put(target, action);
+        return new Subscription(handlers, action, target);
     }
 
     public void remove(String name) {

clients/java/signalr/src/main/java/Subscription.java

@@ -0,0 +1,22 @@
+// 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.
+
+import java.util.List;
+
+public class Subscription {
+    private CallbackMap handlers;
+    private ActionBase action;
+    private String target;
+    public Subscription(CallbackMap handlers, ActionBase action, String target) {
+        this.handlers = handlers;
+        this.action = action;
+        this.target = target;
+    }
+
+    public void unsubscribe() {
+        List<ActionBase> actions = this.handlers.get(target);
+        if (actions != null){
+            actions.remove(action);
+        }
+    }
+}

clients/java/signalr/src/test/java/HubConnectionTest.java

@@ -21,7 +21,6 @@ public class HubConnectionTest {
 
     @Test
     public void RegisteringMultipleHandlersAndBothGetTriggered() throws Exception {
-
         AtomicReference<Double> value = new AtomicReference<>(0.0);
         Transport mockTransport = new MockEchoTransport();
         HubConnection hubConnection = new HubConnection("http://example.com", mockTransport);
@@ -63,7 +62,6 @@ public class HubConnectionTest {
 
     @Test
     public void AddAndRemoveHandlerImmediately() throws Exception {
-
         AtomicReference<Double> value = new AtomicReference<>(0.0);
         Transport mockTransport = new MockEchoTransport();
         HubConnection hubConnection = new HubConnection("http://example.com", mockTransport);
@@ -83,7 +81,6 @@ public class HubConnectionTest {
 
     @Test
     public void RemovingMultipleHandlersWithOneCallToRemove() throws Exception {
-
         AtomicReference<Double> value = new AtomicReference<>(0.0);
         Transport mockTransport = new MockEchoTransport();
         HubConnection hubConnection = new HubConnection("http://example.com", mockTransport);
@@ -107,6 +104,95 @@ public class HubConnectionTest {
         assertEquals(3, value.get(), 0);
     }
 
+    @Test
+    public void RemoveHandlerWithUnsubscribe() throws Exception {
+        AtomicReference<Double> value = new AtomicReference<>(0.0);
+        Transport mockTransport = new MockEchoTransport();
+        HubConnection hubConnection = new HubConnection("http://example.com", mockTransport);
+        Action action = () -> value.getAndUpdate((val) -> val + 1);
+
+        Subscription subscription = hubConnection.on("inc", action);
+
+        assertEquals(0.0, value.get(), 0);
+
+        hubConnection.start();
+        hubConnection.send("inc");
+
+        // Confirming that our handler was called and that the counter property was incremented.
+        assertEquals(1, value.get(), 0);
+
+        subscription.unsubscribe();
+        hubConnection.send("inc");
+        assertEquals(1, value.get(), 0);
+    }
+
+    @Test
+    public void UnsubscribeTwice() throws Exception {
+        AtomicReference<Double> value = new AtomicReference<>(0.0);
+        Transport mockTransport = new MockEchoTransport();
+        HubConnection hubConnection = new HubConnection("http://example.com", mockTransport);
+        Action action = () -> value.getAndUpdate((val) -> val + 1);
+
+        Subscription subscription = hubConnection.on("inc", action);
+
+        assertEquals(0.0, value.get(), 0);
+
+        hubConnection.start();
+        hubConnection.send("inc");
+
+        // Confirming that our handler was called and that the counter property was incremented.
+        assertEquals(1, value.get(), 0);
+
+        subscription.unsubscribe();
+        subscription.unsubscribe();
+        hubConnection.send("inc");
+        assertEquals(1, value.get(), 0);
+    }
+
+    @Test
+    public void RemoveSingleHandlerWithUnsubscribe() throws Exception {
+        AtomicReference<Double> value = new AtomicReference<>(0.0);
+        Transport mockTransport = new MockEchoTransport();
+        HubConnection hubConnection = new HubConnection("http://example.com", mockTransport);
+        Action action = () -> value.getAndUpdate((val) -> val + 1);
+        Action secondAction = () -> value.getAndUpdate((val) -> val + 2);
+
+        Subscription subscription = hubConnection.on("inc", action);
+        Subscription secondSubscription = hubConnection.on("inc", secondAction);
+
+        assertEquals(0.0, value.get(), 0);
+
+        hubConnection.start();
+        hubConnection.send("inc");
+
+        // Confirming that our handler was called and that the counter property was incremented.
+        assertEquals(3, value.get(), 0);
+
+        // This removes the first handler so when "inc" is invoked secondAction should still run.
+        subscription.unsubscribe();
+        hubConnection.send("inc");
+        assertEquals(5, value.get(), 0);
+    }
+
+    @Test
+    public void AddAndRemoveHandlerImmediatelyWithSubscribe() throws Exception {
+        AtomicReference<Double> value = new AtomicReference<>(0.0);
+        Transport mockTransport = new MockEchoTransport();
+        HubConnection hubConnection = new HubConnection("http://example.com", mockTransport);
+        Action action = () -> value.getAndUpdate((val) -> val + 1);
+
+        Subscription sub = hubConnection.on("inc", action);
+        sub.unsubscribe();
+
+        assertEquals(0.0, value.get(), 0);
+
+        hubConnection.start();
+        hubConnection.send("inc");
+
+        // Confirming that the handler was removed.
+        assertEquals(0, value.get(), 0);
+    }
+
     @Test
     public void RegisteringMultipleHandlersThatTakeParamsAndBothGetTriggered() throws Exception {
         AtomicReference<Double> value = new AtomicReference<>(0.0);
@@ -129,7 +215,6 @@ public class HubConnectionTest {
     // We're using AtomicReference<Double> in the send tests instead of int here because Gson has trouble deserializing to Integer
     @Test
     public void SendWithNoParamsTriggersOnHandler() throws Exception {
-
         AtomicReference<Double> value = new AtomicReference<Double>(0.0);
         Transport mockTransport = new MockEchoTransport();
         HubConnection hubConnection = new HubConnection("http://example.com", mockTransport);