6步写好Websocket项目

最近在搞一个AI项目，有许多实时聊天消息，用了websocket长连接。 关于websocket的使用已经不是什么新东西：

Step 1

let ws = new WebSocket("ws://server")

let message = ''
ws.onmessage = (event) => {
    message += event.data
}

以上示例一个最简单的ws连接，当server发送消息后，在onmessage方法内追加消息。

但是大家普遍忽略又十分重要的问题是，websocket使用onmessage作为消息回调（当然也可以addEventListener，不过大同小异），导致它很难把代码写得好看。无论你的程序架构多么好，onmessage回调永远需要拿到外部变量进行处理，此外，当这个情况结合React这类框架时，情况会更加复杂，举个例子：

Step 2

function WebsocketExample() {
  const [message, setMessage] = useState();
  useEffect(() => {
    const ws = new WebSocket('ws://localhost:8080');
    ws.onopen = () => {
      ws.send('Hello!');
    };
    ws.onmessage = (event) => {
      setMessage(message + event.data);
    };
  }, []);

  return ( <>
      {message}
    </> );
}

export default WebsocketExample;

用React写了这样一个最简单的Websocket连接，不断把server端信息追加在message上。这个看起来毫无问题的代码犯了React开发的一个错误，而且难以排查，你会发现错误不是每次必现的，而是时而正常时而错误。

因为onmesasge是一个异步方法，而且是一个持续回调的异步方法，简单的取message变量，在React整个context中，是不能保证message永远是最新的。于是作为一个资深React开发者你做了如下改进

ws.onmessage = (event) => {  
  const newMessage = event.data;  
  setMessages((prevMessages) => `${prevMessage}${newMessage}`);  
};  

至此解决了一个最基本变量问题。

Step 3

实际项目中，业务不可能这么简单，作为一个资深的React开发者，自然会使用各种状态管理框架，于是这样的代码出现了：

import React, { useReducer, useEffect } from 'react';  
  
// 定义 reducer 函数  
const messagesReducer = (state, action) => {  
  switch (action.type) {  
    case 'ADD_MESSAGE':  
      return [...state, action.payload];  
    case 'APPEND_MESSAGE':
      const message = state.find(m => m.id === action.payload.id);
      message.text = message.text + action.payload.text;
      return [...state, message];  
    case 'MODIFY_ROOM_INFO':  
      return [...state, action.payload];  
    case 'VOICE_MESSAGE':  
      return [...state, action.payload];  
    default:  
      throw new Error('Unknown action type');  
  }  
};


ws.onmessage = (event) => {  
  dispatch({ type: event.data.type, payload: event.data.payload });  
};  

大部分搜到的文章，到这一步就停止了，但是我们还要继续

Step 4

到目前为止还算顺利，所有涉及到的业务，都可以比较容易地放在状态中管理，下面增加一些难度了：

1. websocket 向 client 持续推送消息，其中有文字消息也有语音片段
2. 通过消息类型来分别处理，文字就在页面展示，如果是音频就播放
3. 不过要注意的是，音频应该是放在一个队列里播放的，播完一个才能播下一个。

import React, { useReducer, useEffect, useRef } from 'react';  
  
// 定义动作类型  
const ACTION_TYPES = {  
  RECEIVE_TEXT_MESSAGE: 'RECEIVE_TEXT_MESSAGE',  
  RECEIVE_AUDIO_MESSAGE: 'RECEIVE_AUDIO_MESSAGE',  
  PLAY_NEXT_AUDIO: 'PLAY_NEXT_AUDIO',  
  AUDIO_PLAYING_STATUS_CHANGED: 'AUDIO_PLAYING_STATUS_CHANGED',  
};  
  
// 定义初始状态  
const initialState = {  
  textMessages: [],  
  audioMessages: [],  
  currentAudioIndex: -1, // 当前播放的音频在队列中的索引，-1 表示没有正在播放的音频  
  isAudioPlaying: false, // 是否正在播放音频  
};  
  
// 定义 reducer 函数  
function reducer(state, action) {  
  switch (action.type) {  
    case ACTION_TYPES.RECEIVE_TEXT_MESSAGE:  
      return {  
        ...state,  
        textMessages: [...state.textMessages, action.payload],  
      };  
    case ACTION_TYPES.RECEIVE_AUDIO_MESSAGE:  
      return {  
        ...state,  
        audioMessages: [...state.audioMessages, action.payload],  
        // 如果队列为空，则开始播放新接收的音频  
        currentAudioIndex: state.audioMessages.length - 1,  
        isAudioPlaying: state.audioMessages.length === 1,  
      };  
    case ACTION_TYPES.PLAY_NEXT_AUDIO:  
      return {  
        ...state,  
        currentAudioIndex: (state.currentAudioIndex + 1) % state.audioMessages.length,  
        isAudioPlaying: true,  
      };  
    case ACTION_TYPES.AUDIO_PLAYING_STATUS_CHANGED:  
      return {  
        ...state,  
        isAudioPlaying: action.payload,  
      };  
    default:  
      throw new Error('Unknown action type');  
  }  
}  
  
const WebSocketComponent = () => {  
  const [state, dispatch] = useReducer(reducer, initialState);  
  const audioRef = useRef(null);  
  
  useEffect(() => {  
    const ws = new WebSocket('wss://your-websocket-url');  
  
    ws.onmessage = (event) => {  
      const message = JSON.parse(event.data);  
      if (message.type === 'text') {  
        dispatch({ type: ACTION_TYPES.RECEIVE_TEXT_MESSAGE, payload: message.content });  
      } else if (message.type === 'audio') {  
        dispatch({ type: ACTION_TYPES.RECEIVE_AUDIO_MESSAGE, payload: message.content });  
      }  
    };  
  
    ws.onclose = () => {  
      console.log('WebSocket connection closed');  
    };  
  
    ws.onerror = (error) => {  
      console.error('WebSocket error:', error);  
    };  
  
    // 清理函数  
    return () => {  
      ws.close();  
      // 如果组件卸载时有音频正在播放，则停止它  
      if (audioRef.current) {  
        audioRef.current.pause();  
      }  
    };  
  }, []);  
  
  useEffect(() => {  
    // 当有音频需要播放时（即 currentAudioIndex 指向一个有效的音频且当前没有音频正在播放）  
    if (state.currentAudioIndex >= 0 && state.currentAudioIndex < state.audioMessages.length && !state.isAudioPlaying) {  
      const audioContent = state.audioMessages[state.currentAudioIndex];  
      const audioElement = new Audio(`data:audio/mpeg;base64,${audioContent}`);  
      audioRef.current = audioElement; // 存储当前的 Audio 元素  
  
      // 设置音频播放完毕后的回调  
      audioElement.addEventListener('ended', () => {  
        // 播放下一个音频  
        dispatch({ type: ACTION_TYPES.PLAY_NEXT_AUDIO });  
        // 更新播放状态  
        dispatch({ type: ACTION_TYPES.AUDIO_PLAYING_STATUS_CHANGED, payload: false });  
      });  
  
      // 播放音频  
      audioElement.play();  
      // 更新播放状态  
      dispatch({ type: ACTION_TYPES.AUDIO_PLAYING_STATUS_CHANGED, payload: true });  
    }  
  }, [state.currentAudioIndex, state.isAudioPlaying, state.audioMessages]);  
  
  return (  
    <div>  
      <h1>Text Messages</h1>  
      <ul>  
        {state.textMessages.map((msg, index) => (  
          <li key={index}>{msg}</li>  
        ))}  
      </ul>  
      <h1>Audio Messages Queue</h1>  
      <ul>  
        {state.audioMessages.map((msg, index) => (  
          <li key={index}>Audio {index + 1}</li>  
        ))}  
      </ul>  
    </div>  
  );  
};  
  
export default WebSocketComponent;

仅仅这一个功能，已经让代码十分麻烦，主要是因为：

1. onmessage是Websocket的异步回调，本质上它和React状态管理无关，二者需要同步。
2. Audio 对象播放完成的异步回调同理，也需要和React状态管理进行同步。

当越来越多的HTML原生对象的异步回调，需要和React数据进行同步时，会让代码越来越复杂，而且越来越难以拆分，拆分也需要拿到对象的引用，进一步增加代码的复杂性。

Step 5

那么我们要再上难度了：

1. 每个消息有一个id
2. 之前播放的音频，需要可以replay。

作为一个资深React开发者，必然难不倒我们，用一个Map记录每个队列，实时播放，并且还可以replay

import React, { useReducer, useEffect, useRef } from 'react';  
  
// 定义动作类型  
const ACTION_TYPES = {  
  RECEIVE_MESSAGE: 'RECEIVE_MESSAGE', // 接收消息（可以是文本或音频，根据payload区分）  
  PLAY_AUDIO: 'PLAY_AUDIO', // 播放指定ID的音频  
  AUDIO_PLAYING_STATUS_CHANGED: 'AUDIO_PLAYING_STATUS_CHANGED', // 音频播放状态改变  
  RESET_AUDIO_FOR_ID: 'RESET_AUDIO_FOR_ID', // 重置指定ID的音频队列和播放状态  
};  
  
// 定义初始状态  
const initialState = {  
  messages: {}, // 存储所有消息，键是消息ID，值包含type（text/audio）和content（消息内容）  
  audioQueues: {}, // 存储音频队列，键是消息ID，值是音频消息数组  
  isAudioPlaying: {}, // 存储音频播放状态，键是消息ID，值是布尔值  
};  
  
// 定义 reducer 函数  
function reducer(state, action) {  
  switch (action.type) {  
    case ACTION_TYPES.RECEIVE_MESSAGE:  
      const { id, type, content } = action.payload;  
      return {  
        ...state,  
        messages: {  
          ...state.messages,  
          [id]: { type, content },  
        },  
        audioQueues: {  
          ...state.audioQueues,  
          // 如果收到的是音频消息且该ID不存在于audioQueues中，则初始化队列  
          ...(type === 'audio' && !(id in state.audioQueues) ? { [id]: [content] } : {}),  
        },  
      };  
    case ACTION_TYPES.PLAY_AUDIO:  
      const { id } = action.payload;  
      // 检查该ID是否存在音频队列且当前没有正在播放的音频  
      if (id in state.audioQueues && !(id in state.isAudioPlaying || state.isAudioPlaying[id])) {  
        return {  
          ...state,  
          // 设置当前播放的音频ID（这里可以扩展为支持多个音频同时播放的逻辑）  
          // 但为了简化，我们假设一次只能播放一个音频  
          currentPlayingAudioId: id, // 可选，用于跟踪当前播放的音频ID  
          isAudioPlaying: {  
            ...state.isAudioPlaying,  
            [id]: true,  
          },  
        };  
      }  
      return state;  
    case ACTION_TYPES.AUDIO_PLAYING_STATUS_CHANGED:  
      const { id, isPlaying } = action.payload;  
      return {  
        ...state,  
        isAudioPlaying: {  
          ...state.isAudioPlaying,  
          [id]: isPlaying,  
        },  
      };  
    case ACTION_TYPES.RESET_AUDIO_FOR_ID:  
      const { idToReset } = action.payload;  
      return {  
        ...state,  
        audioQueues: {  
          ...state.audioQueues,  
          // 重置指定ID的音频队列  
          [idToReset]: [],  
        },  
        isAudioPlaying: {  
          ...state.isAudioPlaying,  
          // 重置指定ID的播放状态  
          [idToReset]: false,  
        },  
        // 如果当前播放的音频ID是需要重置的ID，则也需要清除它（可选）  
        // currentPlayingAudioId: state.currentPlayingAudioId === idToReset ? null : state.currentPlayingAudioId, // 可选  
      };  
    default:  
      throw new Error('Unknown action type');  
  }  
}  
  
const WebSocketComponent = () => {  
  const [state, dispatch] = useReducer(reducer, initialState);  
  const audioRef = useRef(null);  
  const intervalRef = useRef(null); // 用于定时检查音频播放状态的interval  
  
  useEffect(() => {  
    const ws = new WebSocket('wss://your-websocket-url');  
  
    ws.onmessage = (event) => {  
      const message = JSON.parse(event.data);  
      dispatch({ type: ACTION_TYPES.RECEIVE_MESSAGE, payload: message });  
    };  
  
    ws.onclose = () => {  
      console.log('WebSocket connection closed');  
    };  
  
    ws.onerror = (error) => {  
      console.error('WebSocket error:', error);  
    };  
  
    // 清理函数  
    return () => {  
      ws.close();  
      clearInterval(intervalRef.current); // 清除interval  
      if (audioRef.current) {  
        audioRef.current.pause(); // 停止当前播放的音频  
      }  
    };  
  }, []);  
  
  useEffect(() => {  
    // 检查是否有需要播放的音频且当前没有音频正在播放  
    if (state.currentPlayingAudioId && !(state.currentPlayingAudioId in state.isAudioPlaying || state.isAudioPlaying[state.currentPlayingAudioId])) {  
      const audioQueue = state.audioQueues[state.currentPlayingAudioId];  
      if (audioQueue.length > 0) {  
        const audioContent = audioQueue[0]; // 取队列中的第一个音频消息  
        const audioElement = new Audio(`data:audio/mpeg;base64,${audioContent}`);  
        audioRef.current = audioElement; // 存储当前的Audio元素  
  
        // 设置音频播放完毕后的回调  
        audioElement.addEventListener('ended', () => {  
          // 从队列中移除已播放的音频消息  
          dispatch({  
            type: ACTION_TYPES.RECEIVE_MESSAGE,  
            payload: {  
              id: state.currentPlayingAudioId,  
              type: 'audio', // 假设我们总是从队列中取出音频消息  
              content: audioQueue.slice(1).join(','), // 合并剩余音频消息（这里需要调整，因为content可能不是简单的join就能处理的）  
              // 注意：这里的处理逻辑是简化的，实际上你可能需要更复杂的逻辑来处理音频队列的更新  
              // 比如，你可能需要发送一个新的消息到服务器来更新服务器的音频队列状态  
            },  
          });  
          // 播放下一个音频（如果存在）  
          // 这里我们不再直接dispatch PLAY_AUDIO，因为我们已经通过RECEIVE_MESSAGE更新了队列  
          // 并且下一个useEffect会检查是否需要播放新的音频  
        });  
  
        // 播放音频并更新播放状态  
        audioElement.play();  
        dispatch({  
          type: ACTION_TYPES.AUDIO_PLAYING_STATUS_CHANGED,  
          payload: { id: state.currentPlayingAudioId, isPlaying: true },  
        });  
  
        // 清除之前的interval（如果有的话）并设置一个新的interval来检查音频播放状态  
        clearInterval(intervalRef.current);  
        intervalRef.current = setInterval(() => {  
          // 检查音频是否还在播放（这里只是示例，实际可能不需要这么频繁地检查）  
          // 注意：这个检查可能是不必要的，因为我们已经有了'ended'事件监听器  
          // 但为了演示如何更新状态，我们还是保留了它  
          dispatch({  
            type: ACTION_TYPES.AUDIO_PLAYING_STATUS_CHANGED,  
            payload: { id: state.currentPlayingAudioId, isPlaying: !audioElement.paused },  
          });  
        }, 1000); // 每秒检查一次（这个间隔可以根据需要调整）  
      }  
    }  
  }, [state.currentPlayingAudioId, state.audioQueues, state.isAudioPlaying]); // 注意：这里的依赖数组可能需要根据实际逻辑进行调整  
  
  // 提供一个回放功能的方法（这里以按钮点击为例）  
  const replayAudio = (id) => {  
    // 重置指定ID的音频队列（可选，取决于你的需求）  
    // 如果你不想重置队列，只是从头开始播放，那么可以省略这一步  
    // dispatch({ type: ACTION_TYPES.RESET_AUDIO_FOR_ID, payload: { idToReset: id } });  
  
    // 播放指定ID的第一个音频（这里假设队列不为空）  
    dispatch({ type: ACTION_TYPES.PLAY_AUDIO, payload: { id } });  
  };  
  
  return (  
    <div>  
      <h1>Messages</h1>  
      {Object.keys(state.messages).map((id) => (  
        <div key={id}>  
          <h2>Message ID: {id}</h2>  
          <p>{state.messages[id].type === 'text' ? 'Text: ' + state.messages[id].content : 'Audio Queue (first item shown): ' + (state.audioQueues[id] ? state.audioQueues[id][0].slice(0, 50) + '...' : 'Empty')}</p>  
          {state.messages[id].type === 'audio' && (  
            <button onClick={() => replayAudio(id)}>Replay Audio</button>  
          )}  
        </div>  
      ))}  
    </div>  
  );

Step 6

看到这里，可能大家已经发现了，基于我刚才说的两个原因，仅仅用状态管理框架同步websocket/Audio两个原生对象的异步，已经让代码复杂度急剧攀升。况且这还仅仅是开始，实际业务不会这么简单，举几个例子：

1. server 也不知道语音片段结束，当一个id=123，content='end'消息来时，表示语音片段结束
2. 聊天气泡，需要根据信息状态，展示Server等待、用户输入等待和音频播放状态等
3. 当有语音输入时，需要打断正在播放的音频队列，语音输入结束时，重新播放队列，在此期间推送来的Audio需要忽略。 ....

大家可以挑战一下以上几点需求，在使用任何一种状态管理框架（Jotai/Redux/Recoil/Zustand/Valtio），代码复杂度是什么样的。

我们最开始也是在这些问题上死磕，最后发现没必要，原生对象的异步状态和React异步状态，只要各自处理正确，在合适的时机同步就可以保证业务正常，而Websocket本身也是一种流式数据，用Rxjs就行了。直接上代码：

import { webSocket } from 'rxjs/webSocket';

const subject = webSocket('wss://your-websocket-url');

subject.pipe(
  filter(v => v.type === 'audio'),
  concatMap(async v => {
    return new Promise(resolve => {
      const audio = new Audio(`data:audio/mpeg;base64,${v.data}`);
      audio.onended = () => {
        resolve(v)
      }
      audio.play()
    })
  })
).subscribe({
  next: msg => console.log('message received: ' + msg), // Called whenever there is a message from the server.
  error: err => console.log(err), // Called if at any point WebSocket API signals some kind of error.
  complete: () => console.log('complete') // Called when connection is closed (for whatever reason).
 });

就这么几行，实现了队列顺序播放，而且不需要存储中间变量。

我们基于Rxjs的流式处理能力，完整地改造了websocket消息的处理，让代码简洁、清晰，并且比较好地和状态管理框架进行同步，抛砖引玉，大家可以有什么更好的处理办法也可以交流。