如何让React Context Provider感知上下文消费者挂载状态以按需控制WebSocket连接？

核心思路是让Provider维护一个消费者计数：当计数从0变为1时自动建立WebSocket连接，当计数从1变为0时自动断开连接。以下是两种可落地的实现方式：

方式一：自定义订阅钩子（推荐） #

通过在Context中暴露订阅/取消订阅方法，让消费组件在挂载/卸载时自动上报状态，Provider根据计数变化管理连接。

1. 实现Provider核心逻辑 #

import { createContext, useContext, useState, useEffect, useCallback } from 'react';

const WebSocketContext = createContext({
  state: {},
  subscribe: () => {},
  unsubscribe: () => {},
});

export function WebSocketProvider({ children }) {
  const [wsInstance, setWsInstance] = useState(null);
  const [consumerCount, setConsumerCount] = useState(0);
  const [serverState, setServerState] = useState({});

  // 建立WebSocket连接
  const connect = useCallback(() => {
    if (wsInstance) return;
    const ws = new WebSocket('ws://your-server-endpoint');
    ws.onmessage = (event) => {
      setServerState(prev => ({ ...prev, ...JSON.parse(event.data) }));
    };
    ws.onclose = () => setWsInstance(null);
    setWsInstance(ws);
  }, [wsInstance]);

  // 断开WebSocket连接
  const disconnect = useCallback(() => {
    wsInstance?.close();
    setWsInstance(null);
  }, [wsInstance]);

  // 根据消费者计数自动管理连接
  useEffect(() => {
    if (consumerCount > 0 && !wsInstance) {
      connect();
    } else if (consumerCount === 0 && wsInstance) {
      disconnect();
    }
  }, [consumerCount, wsInstance, connect, disconnect]);

  // 订阅：增加消费者计数
  const subscribe = useCallback(() => {
    setConsumerCount(prev => prev + 1);
  }, []);

  // 取消订阅：减少消费者计数
  const unsubscribe = useCallback(() => {
    setConsumerCount(prev => prev - 1);
  }, []);

  return (
    <WebSocketContext.Provider value={{ state: serverState, subscribe, unsubscribe }}>
      {children}
    </WebSocketContext.Provider>
  );
}

2. 封装消费钩子自动处理订阅 #

创建自定义钩子，让消费组件无需手动调用订阅/取消订阅：

export function useWebSocket() {
  const { state, subscribe, unsubscribe } = useContext(WebSocketContext);
  
  useEffect(() => {
    subscribe();
    return () => unsubscribe();
  }, [subscribe, unsubscribe]);

  return state;
}

此后，任何使用useWebSocket的组件挂载时，都会触发Provider计数增加并建立连接；当最后一个消费组件卸载时，计数归0，自动断开连接。

方式二：Ref跟踪消费者数量（轻量版） #

通过Ref隐式跟踪活跃消费者数量，无需暴露订阅方法，实现更简洁的封装：

import { createContext, useContext, useState, useEffect, useRef, useCallback } from 'react';

const WebSocketContext = createContext({});

export function WebSocketProvider({ children }) {
  const [wsInstance, setWsInstance] = useState(null);
  const [serverState, setServerState] = useState({});
  const consumerCountRef = useRef(0);

  const connect = useCallback(() => {
    const ws = new WebSocket('ws://your-server-endpoint');
    ws.onmessage = (e) => setServerState(prev => ({ ...prev, ...JSON.parse(e.data) }));
    ws.onclose = () => setWsInstance(null);
    setWsInstance(ws);
  }, []);

  const disconnect = useCallback(() => {
    wsInstance?.close();
    setWsInstance(null);
  }, [wsInstance]);

  // 注册消费者
  const register = useCallback(() => {
    consumerCountRef.current += 1;
    if (consumerCountRef.current === 1 && !wsInstance) connect();
  }, [wsInstance, connect]);

  // 注销消费者
  const unregister = useCallback(() => {
    consumerCountRef.current -= 1;
    if (consumerCountRef.current === 0 && wsInstance) disconnect();
  }, [wsInstance, disconnect]);

  return (
    <WebSocketContext.Provider value={{ state: serverState, register, unregister }}>
      {children}
    </WebSocketContext.Provider>
  );
}

// 消费钩子
export function useWebSocket() {
  const { state, register, unregister } = useContext(WebSocketContext);
  
  useEffect(() => {
    register();
    return () => unregister();
  }, [register, unregister]);

  return state;
}

关键注意事项 #

• 避免重复计数：确保useEffect的依赖项正确，防止组件重复渲染时多次触发订阅/取消订阅。
• 重连逻辑：如果需要WebSocket意外断开后自动重连，需在onclose回调中判断当前是否有活跃消费者，仅在有消费者时触发重连。
• 内存泄漏：务必在组件卸载时执行取消订阅逻辑，防止计数异常导致连接无法断开。

内容的提问来源于stack exchange，提问作者Roman Starkov