chore: move uiPort WebSocket into own React hook

packages/trace-viewer/src/ui/uiModeView.tsx

@@ -37,7 +37,7 @@ import { toggleTheme } from '@web/theme';
 import { artifactsFolderName } from '@testIsomorphic/folders';
 import { msToString, settings, useSetting } from '@web/uiUtils';
 import type { ActionTraceEvent } from '@trace/trace';
-import { connect } from './wsPort';
+import { useWebSocket } from './wsPort';
 
 let updateRootSuite: (config: FullConfig, rootSuite: Suite, loadErrors: TestError[], progress: Progress | undefined) => void = () => {};
 let runWatchedTests = (fileNames: string[]) => {};
@@ -84,6 +84,7 @@ export const UIModeView: React.FC<{}> = ({
   const runTestBacklog = React.useRef<Set<string>>(new Set());
   const [collapseAllCount, setCollapseAllCount] = React.useState(0);
   const [isDisconnected, setIsDisconnected] = React.useState(false);
+  const [connectToWebSocketIfNeeded] = useWebSocket();
 
   const inputRef = React.useRef<HTMLInputElement>(null);
 
@@ -99,11 +100,14 @@ export const UIModeView: React.FC<{}> = ({
   React.useEffect(() => {
     inputRef.current?.focus();
     setIsLoading(true);
-    connect({ onEvent: dispatchEvent, onClose: () => setIsDisconnected(true) }).then(send => {
+    connectToWebSocketIfNeeded({
+      onEvent: dispatchEvent,
+      onClose: () => setIsDisconnected(true),
+    }).then(send => {
       sendMessage = send;
       reloadTests();
     });
-  }, [reloadTests]);
+  }, [reloadTests, connectToWebSocketIfNeeded]);
 
   updateRootSuite = React.useCallback((config: FullConfig, rootSuite: Suite, loadErrors: TestError[], newProgress: Progress | undefined) => {
     const selectedProjects = config.configFile ? settings.getObject<string[] | undefined>(config.configFile + ':projects', undefined) : undefined;

packages/trace-viewer/src/ui/workbenchLoader.tsx

@@ -21,7 +21,7 @@ import { MultiTraceModel } from './modelUtil';
 import './workbench.css';
 import { toggleTheme } from '@web/theme';
 import { Workbench } from './workbench';
-import { connect } from './wsPort';
+import { useWebSocket } from './wsPort';
 
 export const WorkbenchLoader: React.FunctionComponent<{
 }> = () => {
@@ -33,6 +33,7 @@ export const WorkbenchLoader: React.FunctionComponent<{
   const [dragOver, setDragOver] = React.useState<boolean>(false);
   const [processingErrorMessage, setProcessingErrorMessage] = React.useState<string | null>(null);
   const [fileForLocalModeError, setFileForLocalModeError] = React.useState<string | null>(null);
+  const [connectToWebSocketIfNeeded] = useWebSocket();
 
   const processTraceFiles = React.useCallback((files: FileList) => {
     const blobUrls = [];
@@ -84,7 +85,7 @@ export const WorkbenchLoader: React.FunctionComponent<{
     }
 
     if (params.has('isServer')) {
-      connect({
+      connectToWebSocketIfNeeded({
         onEvent(method: string, params?: any) {
           if (method === 'loadTrace') {
             setTraceURLs(params!.url ? [params!.url] : []);
@@ -100,7 +101,7 @@ export const WorkbenchLoader: React.FunctionComponent<{
       // Don't re-use blob file URLs on page load (results in Fetch error)
       setTraceURLs(newTraceURLs);
     }
-  }, []);
+  }, [connectToWebSocketIfNeeded]);
 
   React.useEffect(() => {
     (async () => {

packages/trace-viewer/src/ui/wsPort.ts

@@ -14,41 +14,60 @@
  * limitations under the License.
  */
 
-let lastId = 0;
-let _ws: WebSocket;
-const callbacks = new Map<number, { resolve: (arg: any) => void, reject: (arg: Error) => void }>();
+import React from 'react';
 
-export async function connect(options: { onEvent: (method: string, params?: any) => void, onClose: () => void }): Promise<(method: string, params?: any) => Promise<any>> {
-  const guid = new URLSearchParams(window.location.search).get('ws');
-  const ws = new WebSocket(`${window.location.protocol === 'https:' ? 'wss' : 'ws'}://${window.location.hostname}:${window.location.port}/${guid}`);
-  await new Promise(f => ws.addEventListener('open', f));
-  ws.addEventListener('close', options.onClose);
-  ws.addEventListener('message', event => {
-    const message = JSON.parse(event.data);
-    const { id, result, error, method, params } = message;
-    if (id) {
-      const callback = callbacks.get(id);
-      if (!callback)
-        return;
-      callbacks.delete(id);
-      if (error)
-        callback.reject(new Error(error));
-      else
-        callback.resolve(result);
-    } else {
-      options.onEvent(method, params);
-    }
-  });
-  _ws = ws;
-  setInterval(() => sendMessage('ping').catch(() => {}), 30000);
-  return sendMessage;
-}
-
-const sendMessage = async (method: string, params?: any): Promise<any> => {
-  const id = ++lastId;
-  const message = { id, method, params };
-  _ws.send(JSON.stringify(message));
-  return new Promise((resolve, reject) => {
-    callbacks.set(id, { resolve, reject });
-  });
+type ConnectOptions = {
+  onEvent: (method: string, params?: any) => void;
+  onClose: () => void;
 };
+
+type WebSocketMessageSender = (method: string, params?: any) => Promise<any>;
+
+export function useWebSocket(): [(options: ConnectOptions) => Promise<WebSocketMessageSender>] {
+  const lastIdRef = React.useRef(0);
+  const wsRef = React.useRef<WebSocket>();
+  const callbacksRef = React.useRef(new Map<number, { resolve: (arg: any) => void, reject: (arg: Error) => void }>());
+
+  const sendMessage = React.useCallback(async (method: string, params?: any): Promise<any> => {
+    if (!wsRef.current)
+      return;
+    const id = ++lastIdRef.current;
+    const message = { id, method, params };
+    wsRef.current.send(JSON.stringify(message));
+    return new Promise((resolve, reject) => {
+      callbacksRef.current.set(id, { resolve, reject });
+    });
+  }, []);
+
+  const connectIfNeeded = React.useCallback(async (options: ConnectOptions): Promise<WebSocketMessageSender> => {
+    if (wsRef.current)
+      return sendMessage;
+    const guid = new URLSearchParams(window.location.search).get('ws');
+    const ws = new WebSocket(`${window.location.protocol === 'https:' ? 'wss' : 'ws'}://${window.location.hostname}:${window.location.port}/${guid}`);
+    await new Promise(f => ws.addEventListener('open', f));
+    ws.addEventListener('close', options.onClose);
+    ws.addEventListener('message', event => {
+      const message = JSON.parse(event.data);
+      const { id, result, error, method, params } = message;
+      if (id) {
+        const callback = callbacksRef.current.get(id);
+        if (!callback)
+          return;
+        callbacksRef.current.delete(id);
+        if (error)
+          callback.reject(new Error(error));
+        else
+          callback.resolve(result);
+      } else {
+        options.onEvent(method, params);
+      }
+    });
+    wsRef.current = ws;
+    setInterval(() => sendMessage('ping').catch(() => { }), 30000);
+    return sendMessage;
+  }, [sendMessage]);
+
+  return [
+    connectIfNeeded,
+  ];
+}