ReactでWebSocketの接続をカスタムフックにまとめる方法

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

export type WebSocketStatus = 'connecting' | 'open' | 'closed';

export type WebSocketOptions = {
  onMessage?: (event: MessageEvent) => void; // メッセージが届いたときの処理
  reconnect?: boolean; // 思いがけず切れたとき、つなぎ直すか
  retryDelay?: number; // 1回目のつなぎ直しまでの待ち時間（ミリ秒）。失敗のたびに2倍にする
  maxRetryDelay?: number; // 待ち時間の上限（ミリ秒）
};

export function useWebSocket(
  url: string,
  { onMessage, reconnect = true, retryDelay = 500, maxRetryDelay = 8000 }: WebSocketOptions = {},
) {
  const [status, setStatus] = useState<WebSocketStatus>('connecting');
  const socketRef = useRef<WebSocket | null>(null);

  // つなぎ直さずに、届いた時点の最新の onMessage を呼ぶ
  const handleMessage = useEffectEvent((event: MessageEvent) => {
    onMessage?.(event);
  });

  useEffect(() => {
    let socket: WebSocket;
    let timer: ReturnType<typeof setTimeout> | undefined;
    let attempt = 0; // 続けて失敗した回数
    let disposed = false; // 後片付けが済んだか

    function connect() {
      setStatus('connecting');
      socket = new WebSocket(url);
      socketRef.current = socket;
      socket.onopen = () => {
        attempt = 0;
        setStatus('open');
      };
      socket.onmessage = (event) => handleMessage(event);
      socket.onclose = () => {
        if (disposed) return; // 自分で閉じたときは、つなぎ直さない
        socketRef.current = null;
        setStatus('closed');
        if (!reconnect) return;
        const delay = Math.min(retryDelay * 2 ** attempt, maxRetryDelay);
        attempt++;
        timer = setTimeout(connect, delay);
      };
    }

    connect();
    return () => {
      disposed = true;
      clearTimeout(timer);
      socketRef.current = null;
      socket.close();
    };
  }, [url, reconnect, retryDelay, maxRetryDelay]);

  // つながっているときだけ送る。送れなかったら false を返す
  const send = useCallback((data: string) => {
    const socket = socketRef.current;
    if (socket?.readyState !== WebSocket.OPEN) return false;
    socket.send(data);
    return true;
  }, []);

  return { status, send };
}

const { status, send } = useWebSocket(`${baseUrl}/?room=${room}`, {
  onMessage: (event) => setMessages((m) => [...m, String(event.data)]),
});

function handleSend() {
  if (send(text)) {
    setText('');
    setNotice('');
  } else {
    setNotice('つながっていないので、送れませんでした');
  }
}

送ったあとのメッセージ: ["部屋 a に入りました","echo: こんにちは"]
部屋を b にしたあとのメッセージ: ["部屋 a に入りました","echo: こんにちは","部屋 b に入りました"]
サーバーから見た部屋 a の接続: [{"id":1,"open":true,"closed":true,"code":1005}]
サーバーから見た部屋 b の接続: [{"id":2,"open":true,"closed":false,"code":null}]

useEffect(() => {
  const socket = new WebSocket(url);
  socket.onopen = () => setStatus('open');
  socket.onmessage = (event) => onMessage(event);
  socket.onclose = () => setStatus('closed');
  return () => socket.close();
}, [url, onMessage]);

deps: 開いてから2秒のあいだにサーバーが受けた接続 175、そのうち閉じた接続 175、表示「部屋 deps に入りました（入力欄: ）」
deps: 同じ2秒のあいだにページで作った WebSocket 356、そのうち open まで進んだもの 177、open する前に閉じたもの 178、届いたメッセージ 147
effect-event: 開いてから2秒のあいだにサーバーが受けた接続 1、そのうち閉じた接続 0、表示「部屋 deps に入りました（入力欄: ）」

effect-event: 入力欄に abc と打ってからサーバーが送ったあと、サーバーが受けた接続 1、表示「お知らせ（入力欄: abc）」

function connect() {
  socket = new WebSocket(url);
  socket.onopen = () => setStatus('open');
  socket.onclose = () => {
    setStatus('closed');
    setTimeout(connect, 500);
  };
}
connect();
return () => socket.close();

guard: サーバーが受けた接続 1、いま開いている接続 0
no-guard: サーバーが受けた接続 2、いま開いている接続 1

  1回目: ひとつ前が閉じてから 501ms 後に作成、結果 close code 1006
  2回目: ひとつ前が閉じてから 1001ms 後に作成、結果 close code 1006
  3回目: ひとつ前が閉じてから 2001ms 後に作成、結果 close code 1006
  4回目: ひとつ前が閉じてから 4001ms 後に作成、結果 close code 1006
  5回目: ひとつ前が閉じてから 8000ms 後に作成、結果 open
最初に切れた接続の close: {"code":1006,"wasClean":false}
つながったあとにもう一度切ったとき: 閉じてから 501ms 後に作成

素の WebSocket: {"readyStateAtStart":0,"sendWhileConnecting":"InvalidStateError: Failed to execute 'send' on 'WebSocket': Still in CONNECTING state.","readyStateAfterClose":3,"sendAfterClose":"エラーなし","bufferedAmountAfterClose":3}

止まっているあいだに送信を押した結果: {"status":"切断","notice":"つながっていないので、送れませんでした","text":"届く？"}

つなぎ直したあとの入力欄: "届く？"
つなぎ直したあとのメッセージ: ["部屋 a に入りました","部屋 a に入りました","部屋 a に入りました","echo: 届く？"]

StrictMode なし: ページで作った WebSocket [{"opened":true}]、サーバーが受けた接続 1、いま開いている接続 1、画面のメッセージ ["部屋 a に入りました"]
StrictMode あり: ページで作った WebSocket [{"opened":false,"code":1006},{"opened":true}]、サーバーが受けた接続 2、いま開いている接続 1、画面のメッセージ ["部屋 a に入りました"]