如何在React中用useEffect正确管理WebSocket连接，解决泄漏与重复消息？

问题背景 #

我正在开发一款React应用，需要建立WebSocket连接以接收实时更新，但遇到了组件卸载时连接无法正常关闭的问题。当前的实现方案如下：

import { useEffect, useState } from "react";

const useWebSocket = (url: string) => {
  const [messages, setMessages] = useState<string[]>([]);

  useEffect(() => {
    const socket = new WebSocket(url);

    socket.onmessage = (event) => {
      setMessages((prev) => [...prev, event.data]);
    };

    return () => {
      socket.close();
    };
  }, [url]);

  return messages;
};

export default useWebSocket;

存在的问题：

• 有时在页面跳转后，WebSocket连接仍处于开放状态；
• 部分场景下，重新打开同一页面会收到重复消息。

问题原因分析 #

1. 连接关闭不彻底：直接调用socket.close()时，WebSocket可能还处于CONNECTING状态，此时关闭操作不会生效；另外未清理消息监听函数，可能导致组件卸载后仍有回调执行，引发内存泄漏。
2. 重复消息问题：旧连接未关闭时新连接又建立，同一个消息会被多个连接接收；或者连接重建时服务器重发历史消息，客户端没有做去重处理。

改进后的实现方案 #

下面是修复后的自定义Hook，补充了状态管理、监听清理、连接状态判断等关键逻辑：

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

const useWebSocket = (url: string) => {
  const [messages, setMessages] = useState<string[]>([]);
  const socketRef = useRef<WebSocket | null>(null);
  const messageIdsRef = useRef<Set<string>>(new Set()); // 用于消息去重

  useEffect(() => {
    // 仅在socket未初始化时创建连接
    if (!socketRef.current) {
      const socket = new WebSocket(url);
      socketRef.current = socket;

      // 监听连接打开事件
      socket.onopen = () => {
        console.log("WebSocket连接已建立");
      };

      // 处理消息，添加去重逻辑（根据实际消息格式调整）
      socket.onmessage = (event) => {
        try {
          const messageData = JSON.parse(event.data);
          // 假设消息包含唯一id，避免重复接收
          if (!messageIdsRef.current.has(messageData.id)) {
            messageIdsRef.current.add(messageData.id);
            setMessages(prev => [...prev, messageData.content]);
          }
        } catch (e) {
          // 纯文本消息的去重处理
          if (!messageIdsRef.current.has(event.data)) {
            messageIdsRef.current.add(event.data);
            setMessages(prev => [...prev, event.data]);
          }
        }
      };

      // 监听连接关闭，重置实例引用
      socket.onclose = () => {
        console.log("WebSocket连接已关闭");
        socketRef.current = null;
      };

      // 监听错误，主动关闭无效连接
      socket.onerror = (error) => {
        console.error("WebSocket错误:", error);
        socket.close();
        socketRef.current = null;
      };
    }

    // 清理函数：确保连接彻底关闭并清除监听
    return () => {
      if (socketRef.current) {
        // 先移除所有事件监听，避免内存泄漏
        socketRef.current.onopen = null;
        socketRef.current.onmessage = null;
        socketRef.current.onclose = null;
        socketRef.current.onerror = null;
        
        // 仅在连接处于OPEN状态时执行关闭
        if (socketRef.current.readyState === WebSocket.OPEN) {
          socketRef.current.close(1000, "组件卸载，主动关闭连接");
        }
        socketRef.current = null;
      }
      // 重置去重集合，避免重新连接时逻辑失效
      messageIdsRef.current.clear();
    };
  }, [url]);

  return messages;
};

export default useWebSocket;

关键优化点说明 #

• 用useRef存储socket实例：避免每次渲染重新创建连接，同时确保清理函数能访问到最新的socket对象。
• 判断连接状态再关闭：只对OPEN状态的连接执行关闭操作，避免无效调用。
• 清理所有事件监听：关闭连接前移除所有回调，防止组件卸载后异步回调继续执行。
• 消息去重机制：通过useRef维护唯一标识集合，过滤重复消息。
• 错误处理逻辑：出现异常时主动关闭连接并重置实例，避免无效连接占用资源。

最佳实践建议 #

1. 连接复用：如果多个组件需要同一WebSocket连接，可将实例提升到Context或全局状态中，避免重复创建。
2. 自动重连：添加重连逻辑，当连接意外断开时，通过定时器尝试重建连接（注意设置重连间隔上限，避免频繁请求）。
3. 标准化消息格式：建议服务端和客户端用JSON传递消息，包含唯一id字段，方便去重和追踪。
4. 全量清理异步操作：组件卸载时，除了WebSocket，还要清理定时器、订阅等所有异步任务。
5. 暴露连接状态：在Hook中添加isConnected状态，方便组件根据连接状态展示加载、断开提示等UI。

内容的提问来源于stack exchange，提问作者Sandro Tushurashvili