Acquire HubConnectionStateLock before Send/Invoke/Stream

src/SignalR/clients/java/signalr/src/main/java/com/microsoft/signalr/HubConnection.java

@@ -57,6 +57,7 @@ public class HubConnection {
     private Map<String, Observable> streamMap = new ConcurrentHashMap<>();
     private TransportEnum transportEnum = TransportEnum.ALL;
     private String connectionId;
+    private final Lock connectionStateLock = new ReentrantLock();
     private final Logger logger = LoggerFactory.getLogger(HubConnection.class);
 
     /**
@@ -505,9 +506,15 @@ private void stopConnection(String errorMessage) {
                 exception = new RuntimeException(errorMessage);
                 logger.error("HubConnection disconnected with an error {}.", errorMessage);
             }
-            if (connectionState != null) {
-                connectionState.cancelOutstandingInvocations(exception);
-                connectionState = null;
+
+            connectionStateLock.lock();
+            try {
+                if (connectionState != null) {
+                    connectionState.cancelOutstandingInvocations(exception);
+                    connectionState = null;
+                }
+            } finally {
+                connectionStateLock.unlock();
             }
 
             logger.info("HubConnection stopped.");
@@ -537,8 +544,13 @@ private void stopConnection(String errorMessage) {
      * @param args   The arguments to be passed to the method.
      */
     public void send(String method, Object... args) {
-        if (hubConnectionState != HubConnectionState.CONNECTED) {
-            throw new RuntimeException("The 'send' method cannot be called if the connection is not active.");
+        hubConnectionStateLock.lock();
+        try {
+            if (hubConnectionState != HubConnectionState.CONNECTED) {
+                throw new RuntimeException("The 'send' method cannot be called if the connection is not active.");
+            }
+        } finally {
+            hubConnectionStateLock.unlock();
         }
 
         sendInvocationMessage(method, args);
@@ -638,15 +650,23 @@ public Completable invoke(String method, Object... args) {
      */
     @SuppressWarnings("unchecked")
     public <T> Single<T> invoke(Class<T> returnType, String method, Object... args) {
-        if (hubConnectionState != HubConnectionState.CONNECTED) {
-            throw new RuntimeException("The 'invoke' method cannot be called if the connection is not active.");
-        }
+        ConnectionState localConnectionState;
+        InvocationRequest irq;
+        String id;
+        hubConnectionStateLock.lock();
+        try {
+            if (hubConnectionState != HubConnectionState.CONNECTED) {
+                throw new RuntimeException("The 'invoke' method cannot be called if the connection is not active.");
+            }
 
-        String id = connectionState.getNextInvocationId();
+            id = connectionState.getNextInvocationId();
+            irq = new InvocationRequest(returnType, id);
+            connectionState.addInvocation(irq);
+        } finally {
+            hubConnectionStateLock.unlock();
+        }
 
         SingleSubject<T> subject = SingleSubject.create();
-        InvocationRequest irq = new InvocationRequest(returnType, id);
-        connectionState.addInvocation(irq);
 
         // forward the invocation result or error to the user
         // run continuations on a separate thread
@@ -677,12 +697,23 @@ public <T> Single<T> invoke(Class<T> returnType, String method, Object... args)
      */
     @SuppressWarnings("unchecked")
     public <T> Observable<T> stream(Class<T> returnType, String method, Object ... args) {
-        String invocationId = connectionState.getNextInvocationId();
+        String invocationId;
+        InvocationRequest irq;
+        hubConnectionStateLock.lock();
+        try {
+            if (hubConnectionState != HubConnectionState.CONNECTED) {
+                throw new RuntimeException("The 'stream' method cannot be called if the connection is not active.");
+            }
+
+            invocationId = connectionState.getNextInvocationId();
+            irq = new InvocationRequest(returnType, invocationId);
+            connectionState.addInvocation(irq);
+        } finally {
+            hubConnectionStateLock.unlock();
+        }
+
         AtomicInteger subscriptionCount = new AtomicInteger();
-        InvocationRequest irq = new InvocationRequest(returnType, invocationId);
-        connectionState.addInvocation(irq);
         ReplaySubject<T> subject = ReplaySubject.create();
-
         Subject<Object> pendingCall = irq.getPendingCall();
         pendingCall.subscribe(result -> {
             // Primitive types can't be cast with the Class cast function
@@ -700,7 +731,14 @@ public <T> Observable<T> stream(Class<T> returnType, String method, Object ... a
             if (subscriptionCount.decrementAndGet() == 0) {
                 CancelInvocationMessage cancelInvocationMessage = new CancelInvocationMessage(invocationId);
                 sendHubMessage(cancelInvocationMessage);
-                connectionState.tryRemoveInvocation(invocationId);
+                connectionStateLock.lock();
+                try {
+                    if (connectionState != null) {
+                        connectionState.tryRemoveInvocation(invocationId);
+                    }
+                } finally {
+                    connectionStateLock.unlock();
+                }
                 subject.onComplete();
             }
         });

src/SignalR/clients/java/signalr/src/test/java/com/microsoft/signalr/HubConnectionTest.java

@@ -1608,6 +1608,15 @@ public void cannotInvokeBeforeStart()  {
         assertEquals("The 'invoke' method cannot be called if the connection is not active.", exception.getMessage());
     }
 
+    @Test
+    public void cannotStreamBeforeStart()  {
+        HubConnection hubConnection = TestUtils.createHubConnection("http://example.com");
+        assertEquals(HubConnectionState.DISCONNECTED, hubConnection.getConnectionState());
+
+        Throwable exception = assertThrows(RuntimeException.class, () -> hubConnection.stream(String.class, "inc", "arg1"));
+        assertEquals("The 'stream' method cannot be called if the connection is not active.", exception.getMessage());
+    }
+
     @Test
     public void doesNotErrorWhenReceivingInvokeWithIncorrectArgumentLength()  {
         MockTransport mockTransport = new MockTransport();
@@ -2036,7 +2045,7 @@ public void authorizationHeaderFromNegotiateGetsSetToNewValue() {
 
         TestHttpClient client = new TestHttpClient()
                 .on("POST", "http://example.com/negotiate", (req) -> {
-                    if(redirectCount.get() == 0){
+                    if (redirectCount.get() == 0) {
                         redirectCount.incrementAndGet();
                         redirectToken.set(req.getHeaders().get("Authorization"));
                         return Single.just(new HttpResponse(200, "", "{\"url\":\"http://testexample.com/\",\"accessToken\":\"firstRedirectToken\"}"));

src/SignalR/clients/java/signalr/src/test/java/com/microsoft/signalr/WebSocketTransportUrlFormatTest.java

@@ -8,7 +8,6 @@
 import java.util.HashMap;
 import java.util.stream.Stream;
 
-import io.reactivex.Single;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.Arguments;
 import org.junit.jupiter.params.provider.MethodSource;

src/SignalR/clients/java/signalr/src/test/java/com/microsoft/signalr/sample/Chat.java

@@ -37,6 +37,6 @@ public static void main(String[] args) {
             hubConnection.send("Send", message);
         }
 
-        hubConnection.stop();
+        hubConnection.stop().blockingAwait();
     }
 }