Fix never ending socket connection in streaming

AppiumTestDistribution · Apr 11, 2024 · 0019277 · 0019277
1 parent 53d4265
commit 0019277
Show file tree

Hide file tree

Showing 9 changed files with 112 additions and 95 deletions.
diff --git a/src/modules b/src/modules
diff --git a/src/plugin.ts b/src/plugin.ts
@@ -408,7 +408,7 @@ class DevicePlugin extends BasePlugin {
       if (device.platform.toLowerCase() === 'ios' && !isRemoteOrCloudSession) {
         // eslint-disable-next-line @typescript-eslint/ban-ts-comment
         // @ts-ignore
-        const { server, port, scheme, sessionId } = Object.values(driver.sessions)[0].wda.jwproxy;
+        const { server, port, scheme } = driver.sessions[sessionId].wda.jwproxy;
 
         const proxiedInfo = {
           proxyUrl: `${scheme}://${server}:${port}`,
@@ -577,7 +577,12 @@ class DevicePlugin extends BasePlugin {
     await unblockDeviceMatchingFilter({ session_id: sessionId });
     log.info(`📱 Unblocking the device that is blocked for session ${sessionId}`);
     const res = await next();
-    await registerAndroidWebSocketHandlers(DevicePlugin.httpServer, DevicePlugin.adbInstance);
+    try {
+      await waitForWebsocketToBeDeregister(this.pluginArgs, DevicePlugin.httpServer);
+      await DevicePlugin.registerWebSocket(this.pluginArgs);
+    } catch (err) {
+      log.info('Socket server not removed within 5000ms. So not registering again');
+    }
     return res;
   }
 }

diff --git a/web/package-lock.json b/web/package-lock.json
diff --git a/web/package.json b/web/package.json
@@ -20,7 +20,8 @@
     "react-chartjs-2": "^5.2.0",
     "react-dom": "^18.2.0",
     "react-json-view": "^1.21.3",
-    "react-router-dom": "^6.21.1"
+    "react-router-dom": "^6.21.1",
+    "react-use-websocket": "^3.0.0"
   },
   "devDependencies": {
     "@types/react": "^18.2.43",

diff --git a/web/src/assets/device-loading.gif b/web/src/assets/device-loading.gif
diff --git a/web/src/components/streaming/AndroidStream.tsx b/web/src/components/streaming/AndroidStream.tsx
@@ -6,37 +6,27 @@ import DeviceFarmApiService from '../../api-service';
 import HomeIcon from '@mui/icons-material/Home';
 import { Camera, Upload, Close } from '@mui/icons-material';
 import { toolBarControl } from './util.ts';
+import DeviceLoading from '../../assets/device-loading.gif';
+import useWebSocket from 'react-use-websocket';
 
 const MAX_HEIGHT = 720;
 const MAX_WIDTH = 720;
 
 function AndroidStream() {
-  const [imageSrc, setImageSrc] = useState('');
+  const [imageSrc, setImageSrc] = useState(DeviceLoading);
   const [file, setFile] = useState(null);
 
   const containerElement = useRef<HTMLDivElement>(null);
   const videoElement = useRef<HTMLImageElement>(null);
   const canvasElement = useRef<HTMLCanvasElement>(null);
   let interactionHandler: SimpleInterationHandler | null;
 
-  // eslint-disable-next-line prefer-const
-  let [ws, setWebSocket] = useState<WebSocket | undefined>(undefined);
   const handleWebSocketMessage = (event: { data: any }) => {
     const blob = event.data;
     const url = URL.createObjectURL(blob);
     setImageSrc(url);
   };
 
-  const createWebSocketConnection = (wsUrl: string) => {
-    ws = new WebSocket(wsUrl);
-    ws.addEventListener('message', handleWebSocketMessage);
-    ws.addEventListener('close', () => {
-      console.log('WebSocket connection closed. Reconnecting...');
-      createWebSocketConnection(wsUrl);
-    });
-    return ws;
-  };
-
   const getParamsFromUrl = () => {
     if (window.location.hash.includes('?')) {
       const params = new URLSearchParams(window.location.hash.split('?')[1]);
@@ -51,22 +41,15 @@ function AndroidStream() {
       return { port: 8004, host: '127.0.0.1', udid: '' };
     }
   };
-  useEffect(() => {
-    const { host, udid, port } = getParamsFromUrl() as any;
-    const wsUrl = `ws://${host}:${port}/android-stream/${udid}`;
-
-    createWebSocketConnection(wsUrl);
-    setWebSocket(ws);
-    return () => {
-      // Clean up the WebSocket connection when the component is unmounted
-      // eslint-disable-next-line @typescript-eslint/ban-ts-comment
-      // @ts-expect-error
-      ws.removeEventListener('message', handleWebSocketMessage);
-      // eslint-disable-next-line @typescript-eslint/ban-ts-comment
-      // @ts-expect-error
-      ws.close();
-    };
-  }, []);
+  const { host, udid, port } = getParamsFromUrl() as any;
+  const wsUrl = `ws://${host}:${port}/android-stream/${udid}`;
+  const { sendMessage } = useWebSocket(wsUrl, {
+    share: false,
+    shouldReconnect: () => true,
+    onMessage: handleWebSocketMessage,
+    reconnectInterval: 1500,
+    reconnectAttempts: 15,
+  });
 
   useEffect(() => {
     const { width, height } = getParamsFromUrl() as any;
@@ -82,14 +65,14 @@ function AndroidStream() {
         videoElement.current as any,
         canvasElement.current,
         containerElement.current,
-        ws,
+        { send: sendMessage },
         { width, height },
       );
     }
   }, []);
 
   async function onToolbarControlClick(controlAction: string) {
-    await toolBarControl(ws, controlAction, getParamsFromUrl);
+    await toolBarControl({ send: sendMessage } as any, controlAction, getParamsFromUrl);
   }
   const handleFileChange = (event: React.ChangeEvent<HTMLInputElement>) => {
     // eslint-disable-next-line @typescript-eslint/ban-ts-comment
@@ -185,9 +168,7 @@ function AndroidStream() {
               icon: (
                 <div>
                   <label htmlFor="input-file">Upload APK</label>
-                  <Upload
-                    onClick={uploadFile}
-                  />
+                  <Upload onClick={uploadFile} />
                   <input
                     type="file"
                     onChange={handleFileChange}
@@ -196,7 +177,7 @@ function AndroidStream() {
                   />
                 </div>
               ),
-              name: ''
+              name: '',
             },
             {
               action: 'close',

diff --git a/web/src/components/streaming/ios-stream.tsx b/web/src/components/streaming/ios-stream.tsx
@@ -5,12 +5,38 @@ import { StreamingToolBar } from './toolbar';
 import { SimpleInterationHandler } from '../../libs/simple-interation-handler';
 import { Camera, Close, Upload } from '@mui/icons-material';
 import { toolBarControl, uploadFile } from './util.ts';
+import DeviceLoading from '../../assets/device-loading.gif';
+import useWebSocket from 'react-use-websocket';
 
 const MAX_HEIGHT = 720;
 const MAX_WIDTH = 720;
 
 function IOSStream() {
+  const getParamsFromUrl = () => {
+    if (window.location.hash.includes('?')) {
+      const params = new URLSearchParams(window.location.hash.split('?')[1]);
+      return {
+        port: params.get('port'),
+        udid: params.get('udid'),
+        host: params.get('host'),
+        width: params.get('width'),
+        height: params.get('height'),
+        streamPort: params.get('streamPort'),
+      };
+    } else {
+      return { port: 8004, host: '127.0.0.1', udid: '' };
+    }
+  };
+
   // const [imageSrc, setImageSrc] = useState('');
+  const { host, port, udid, width, height, streamPort } = getParamsFromUrl() as any;
+  const wsUrl = `ws://${host}:${port}/ios-stream/${udid}`;
+  const { sendMessage } = useWebSocket(wsUrl, {
+    share: false,
+    shouldReconnect: () => true,
+    reconnectInterval: 1500,
+    reconnectAttempts: 15,
+  });
 
   const containerElement = useRef<HTMLDivElement>(null);
   const videoElement = useRef<HTMLImageElement>(null);
@@ -19,22 +45,10 @@ function IOSStream() {
   let interactionHandler: SimpleInterationHandler | null;
 
   // eslint-disable-next-line prefer-const
-  let [ws, setWebSocket] = useState<WebSocket | undefined>(undefined);
-
-  const createWebSocketConnection = (wsUrl: string) => {
-    ws = new WebSocket(wsUrl);
-    // ws.addEventListener('message', handleWebSocketMessage);
-    ws.addEventListener('close', () => {
-      console.log('WebSocket connection closed. Reconnecting...');
-      createWebSocketConnection(wsUrl);
-    });
-    setWebSocket(ws);
-    return ws;
-  };
 
   const uploadAUT = async () => {
-    await uploadFile(file, getParamsFromUrl)
-  }
+    await uploadFile(file, getParamsFromUrl);
+  };
 
   const handleFileChange = (event: React.ChangeEvent<HTMLInputElement>) => {
     // eslint-disable-next-line @typescript-eslint/ban-ts-comment
@@ -45,36 +59,7 @@ function IOSStream() {
     setFile(selectedFile);
   };
 
-  const getParamsFromUrl = () => {
-    if (window.location.hash.includes('?')) {
-      const params = new URLSearchParams(window.location.hash.split('?')[1]);
-      return {
-        port: params.get('port'),
-        udid: params.get('udid'),
-        host: params.get('host'),
-        width: params.get('width'),
-        height: params.get('height'),
-        streamPort: params.get('streamPort'),
-      };
-    } else {
-      return { port: 8004, host: '127.0.0.1', udid: '' };
-    }
-  };
-  useEffect(() => {
-    const { host, port, udid } = getParamsFromUrl() as any;
-    const wsUrl = `ws://${host}:${port}/ios-stream/${udid}`;
-
-    createWebSocketConnection(wsUrl);
-    return () => {
-      // eslint-disable-next-line @typescript-eslint/ban-ts-comment
-      // @ts-expect-error
-      ws.close();
-    };
-  }, []);
-
   useEffect(() => {
-    const { width, height } = getParamsFromUrl() as any;
-
     if (
       !interactionHandler &&
       canvasElement.current &&
@@ -86,14 +71,20 @@ function IOSStream() {
         videoElement.current as any,
         canvasElement.current,
         containerElement.current,
-        ws,
+        { send: sendMessage } as any,
         { width, height },
       );
     }
   }, []);
-  const { host, port, streamPort } = getParamsFromUrl() as any;
+
   async function onToolbarControlClick(controlAction: string) {
-    await toolBarControl(ws, controlAction, getParamsFromUrl);
+    await toolBarControl(
+      {
+        send: sendMessage,
+      } as any,
+      controlAction,
+      getParamsFromUrl,
+    );
   }
   return (
     <div className="streaming-container">
@@ -108,10 +99,14 @@ function IOSStream() {
         >
           <img
             style={{
-              maxHeight: 730 + 'px',
-              maxWidth: 730 + 'px',
+              maxHeight: MAX_HEIGHT + 'px',
+              maxWidth: MAX_WIDTH + 'px',
               width: 'auto',
               position: 'absolute',
+              backgroundImage: `url(${DeviceLoading})`,
+              backgroundRepeat: 'no-repeat',
+              backgroundAttachment: 'fixed',
+              backgroundPosition: 'center',
             }}
             src={`http://${host}:${port}/device-farm/api/dashboard/mjpeg-stream/${streamPort}`}
             ref={videoElement}
@@ -137,9 +132,7 @@ function IOSStream() {
               icon: (
                 <div>
                   <label htmlFor="input-file">Upload IPA/App</label>
-                  <Upload
-                    onClick={uploadAUT}
-                  />
+                  <Upload onClick={uploadAUT} />
                   <input
                     type="file"
                     onChange={handleFileChange}
@@ -148,7 +141,7 @@ function IOSStream() {
                   />
                 </div>
               ),
-              name: ''
+              name: '',
             },
             {
               action: 'close',

diff --git a/web/src/components/streaming/streaming.css b/web/src/components/streaming/streaming.css
@@ -9,7 +9,7 @@
 .device-container {
   display: flex;
   flex-direction: row;
-  gap: 0;
+  gap: 10px;
 }
 
 .toolbar-container {