Academy
hardExtra#97

WebSocket connection hook

Prompt

Create a custom useWebSocket hook that:

  1. Connects to a WebSocket URL
  2. Returns the latest message
  3. Reconnects on disconnect — but never after unmount
  4. Cleans up the socket AND any pending reconnect timer on unmount

Solution

function useWebSocket(url) {
  const [lastMessage, setLastMessage] = useState(null)
  const wsRef = useRef(null)
  const timerRef = useRef(null)
  const closedRef = useRef(false)

  const connect = useCallback(() => {
    wsRef.current = new WebSocket(url)
    wsRef.current.onmessage = (e) => setLastMessage(e.data)
    wsRef.current.onclose = () => {
      // reconnect only while mounted — cleanup's close() also fires onclose
      if (!closedRef.current) timerRef.current = setTimeout(connect, 3000)
    }
  }, [url])

  useEffect(() => {
    closedRef.current = false
    connect()
    return () => {
      closedRef.current = true
      clearTimeout(timerRef.current)
      wsRef.current?.close()
    }
  }, [connect])

  const sendMessage = useCallback((msg) => wsRef.current?.send(msg), [])
  return { lastMessage, sendMessage }
}
Mentor's take

A socket is imperative, stateful, and outlives renders — exactly what refs are for. The hook's whole job is mapping that lifecycle onto React's: wsRef holds the instance (a socket is not render data, so it must not live in state), and only lastMessage is state, because it's the one thing that should trigger a re-render.

The part that separates candidates is the unmount interaction. Cleanup calls close() — which fires onclose — and a naive onclose = () => setTimeout(connect, 3000) then schedules a reconnect for a component that no longer exists: a zombie socket that reconnects forever, plus a setState-after-unmount warning. The closedRef guard breaks that loop, and clearTimeout covers the other race — unmounting while a reconnect is already pending. Works-in-dev, leaks-in-production is exactly this shape of bug.

Production hardening to name unprompted: exponential backoff with jitter instead of a fixed 3s (a fleet of clients reconnecting in lockstep after a server blip is a thundering herd), application-level heartbeat ping/pong because proxies and load balancers silently kill idle connections, and re-authenticate plus re-subscribe on every reconnect — the new socket knows nothing about the old one's session.

Red flag: creating the socket in the render body (one socket per render) or storing it in useState — both confuse "instance I manage" with "data I render."

Say it: "The socket lives in a ref, the latest message in state, and my cleanup both closes the socket and disarms the reconnect — in production I'd add backoff with jitter and heartbeats, because intermediaries kill idle connections."