react hook websocket 封装

websocket 封装

新建 websocket.ts 文件

import {
    
     useState, useRef, useEffect } from 'react';

const useWebsocket = ({
     
      url }) => {
    
    
  const ws = useRef<WebSocket | null>(null);
  // socket 数据
  const [wsData, setMessage] = useState({
    
    });
  //  socket 状态
  const [readyState, setReadyState] = useState<any>({
    
     key: 0, value: '正在连接中' });

  const creatWebSocket = () => {
    
    
    const stateArr = [
      {
    
     key: 0, value: '正在连接中' },
      {
    
     key: 1, value: '已经连接并且可以通讯' },
      {
    
     key: 2, value: '连接正在关闭' },
      {
    
     key: 3, value: '连接已关闭或者没有连接成功' },
    ];

    try {
    
    
      ws.current = new WebSocket(url);

      ws.current.onopen = () => {
    
    
        setReadyState(stateArr[ws.current?.readyState ?? 0]);
      };
      ws.current.onclose = () => {
    
    
        setReadyState(stateArr[ws.current?.readyState ?? 0]);
      };
      ws.current.onerror = () => {
    
    
        setReadyState(stateArr[ws.current?.readyState ?? 0]);
      };
      ws.current.onmessage = e => {
    
    
         const {
    
     data, type } = (...JSON.parse(e.data)) || {
    
    };
        switch (
          type // type 是跟后端约定的
        ) {
    
    
          case '101':
            setMessage({
    
     ...JSON.parse(e.data) },review: data); //根据自身情况进行修改
            break;
          case '102':
            setMessage({
    
     ...JSON.parse(e.data) },pipelineResults: data);//根据自身情况进行修改
            break;
          default:
            setMessage({
    
     ...JSON.parse(e.data), ...data });//根据自身情况进行修改
            break;
        }
      };
    } catch (error) {
    
    
      console.log(error);
    }
  };

  const webSocketInit = () => {
    
    
    if (!ws.current || ws.current.readyState === 3) {
    
    
      creatWebSocket();
    }
  };

  //  关闭 WebSocket
  const closeWebSocket = () => {
    
    
    ws.current?.close();
  };

  // 发送数据
  const sendMessage = (str: string) => {
    
    
    console.log(str, '------------str');
    ws.current?.send(str);
  };

  //重连
  const reconnect = () => {
    
    
    try {
    
    
      closeWebSocket();
      ws.current = null;
      creatWebSocket();
    } catch (e) {
    
    
      console.log(e);
    }
  };

  useEffect(() => {
    
    
    webSocketInit();
    return () => {
    
    
      ws.current?.close();
    };
  }, [ws]);

  //    wsData （获得的 socket 数据）、readyState（当前 socket 状态）、closeWebSocket （关闭 socket）、reconnect（重连）
  return {
    
    
    wsData,
    readyState,
    closeWebSocket,
    reconnect,
    sendMessage,
  };
};
export default useWebsocket;

websocket 使用

import useWebsocket from '@/utils/webSocket';

const [isLocalPage, setIsLocalPage] = useState(true);
const {
    
     wsData, readyState, verify, closeWebSocket, reconnect, sendMessage } = useWebsocket({
    
    
    url: 'ws://XXX.XXX.XXX.XXX:9002/ws/test/1', // 此参数为websocket地址
  });
  useEffect(() => {
    
    
    
    // 主动请求
    if (readyState.key === 1 && isLocalPage) {
    
    
    console.log(readyState, wsData, '----------readyState');
      sendMessage('hello websocket');//发送信息
    }

    // 接受到socket数据， 进行业务逻辑处理
    if (Object.keys(wsData).length !== 0) {
    
    
      console.log(wsData);
    }

    // 如果是已关闭且是当前页面自动重连
    if (readyState.key === 3 && isLocalPage) {
    
    
      reconnect();
    }
    // 不是当前页面 清空 webSocket 此处为优化代码使用的，不需要可以直接删除。
    if (!isLocalPage) {
    
    
      closeWebSocket();
    }
  }, [wsData, readyState, isLocalPage]);
 /*
   ** 判断用户是否离开当前页面，离开后不请求轮询接口，回到当前页面重新执行轮询
   */
  const checkIsLocalPage = () => {
    
    
    document.addEventListener('visibilitychange', function () {
    
    
      // 页面变为不可见时触发
      if (document.visibilityState === 'hidden') {
    
    
        setIsLocalPage(false)
      }
      // 页面变为可见时触发
      if (document.visibilityState === 'visible') {
    
    
        setIsLocalPage(true)
      }
    })
  }
  
useEffect(() => {
    
    
  checkIsLocalPage()
})