[TS Client] Properly wait for WebSocket close (#2804)

clients/ts/signalr/src/HttpConnection.ts

@@ -317,7 +317,7 @@ export class HttpConnection implements IConnection {
         return false;
     }
 
-    private async stopConnection(error?: Error): Promise<void> {
+    private stopConnection(error?: Error): void {
         this.transport = undefined;
 
         // If we have a stopError, it takes precedence over the error from the transport

clients/ts/signalr/src/WebSocketTransport.ts

@@ -68,17 +68,7 @@ export class WebSocketTransport implements ITransport {
                 }
             };
 
-            webSocket.onclose = (event: CloseEvent) => {
+            webSocket.onclose = (event: CloseEvent) => this.close(event);
-                // webSocket will be null if the transport did not start successfully
-                this.logger.log(LogLevel.Trace, "(WebSockets transport) socket closed.");
-                if (this.onclose) {
-                    if (event.wasClean === false || event.code !== 1000) {
-                        this.onclose(new Error(`Websocket closed with status code: ${event.code} (${event.reason})`));
-                    } else {
-                        this.onclose();
-                    }
-                }
-            };
         });
     }
 
@@ -94,9 +84,30 @@ export class WebSocketTransport implements ITransport {
 
     public stop(): Promise<void> {
         if (this.webSocket) {
+            // Clear websocket handlers because we are considering the socket closed now
+            this.webSocket.onclose = () => {};
+            this.webSocket.onmessage = () => {};
+            this.webSocket.onerror = () => {};
             this.webSocket.close();
             this.webSocket = undefined;
+
+            // Manually invoke onclose callback inline so we know the HttpConnection was closed properly before returning
+            // This also solves an issue where websocket.onclose could take 18+ seconds to trigger during network disconnects
+            this.close(undefined);
         }
+
         return Promise.resolve();
     }
+
+    private close(event?: CloseEvent): void {
+        // webSocket will be null if the transport did not start successfully
+        this.logger.log(LogLevel.Trace, "(WebSockets transport) socket closed.");
+        if (this.onclose) {
+            if (event && (event.wasClean === false || event.code !== 1000)) {
+                this.onclose(new Error(`Websocket closed with status code: ${event.code} (${event.reason})`));
+            } else {
+                this.onclose();
+            }
+        }
+    }
 }