score:15

Accepted answer

I'm not convinced that an effect is the best place for this. If it's application-level, it may be simpler to implement it in its own module, and bring that in, where needed.

Nevertheless, to get this to work, you should consider that you're managing two separate lifecycles: the component lifecycle, and the websocket lifecycle. To make it work as you want, you have to ensure that each state change in one aligns with a state change in the other.

First, keep in mind that your effect runs every time the dependencies in the array change. So, in your example, your effect runs every time you set wsState.

The other thing to keep in mind is that your cleanup function is called every time wsState changes, which you're doing twice in your effect (setting it to true on open, and false on close). This means that when you create a new socket, and it fails to connect, the close event fires, and it queues up a state change.

Each time it attempts to connect, it sets wsState to true (which queues a re-run of your effect), tries and fails to connect, finally setting another timeout, which updates the state to false. But, not before the effect runs again, trying to set the state to true, etc.

To fix this, start with the effect lifecycle. When should your effect run? When should it be cleaned up? A few thoughts:

  1. The effect should run once during the first render, but not during subsequent renders
  2. The effect should be cleaned up when the WebSocket disconnects
  3. The effect should be re-run after a timeout, triggering a reconnect

What does this mean for the component? You don't want to include the WS state as a dependency. But, you do need state to trigger it to re-run after the timeout.

Here's what this looks like:

import React, { useRef, useState, useEffect } from 'react';

const URL = 'ws://localhost:8888';

export default function App() {

  const clientRef = useRef(null);
  const [waitingToReconnect, setWaitingToReconnect] = useState(null);
  const [messages, setMessages] = useState([]);
  const [isOpen, setIsOpen] = useState(false);

  function addMessage(message) {
    setMessages([...messages, message]); 
  }

  useEffect(() => {

    if (waitingToReconnect) {
      return;
    }

    // Only set up the websocket once
    if (!clientRef.current) {
      const client = new WebSocket(URL);
      clientRef.current = client;

      window.client = client;

      client.onerror = (e) => console.error(e);

      client.onopen = () => {
        setIsOpen(true);
        console.log('ws opened');
        client.send('ping');
      };

      client.onclose = () => {

        if (clientRef.current) {
          // Connection failed
          console.log('ws closed by server');
        } else {
          // Cleanup initiated from app side, can return here, to not attempt a reconnect
          console.log('ws closed by app component unmount');
          return;
        }

        if (waitingToReconnect) {
          return;
        };

        // Parse event code and log
        setIsOpen(false);
        console.log('ws closed');

        // Setting this will trigger a re-run of the effect,
        // cleaning up the current websocket, but not setting
        // up a new one right away
        setWaitingToReconnect(true);

        // This will trigger another re-run, and because it is false,
        // the socket will be set up again
        setTimeout(() => setWaitingToReconnect(null), 5000);
      };

      client.onmessage = message => {
        console.log('received message', message);
        addMessage(`received '${message.data}'`);
      };


      return () => {

        console.log('Cleanup');
        // Dereference, so it will set up next time
        clientRef.current = null;

        client.close();
      }
    }

  }, [waitingToReconnect]);

  return (
    <div>
      <h1>Websocket {isOpen ? 'Connected' : 'Disconnected'}</h1>
      {waitingToReconnect && <p>Reconnecting momentarily...</p>}
      {messages.map(m => <p>{JSON.stringify(m, null, 2)}</p>)}
    </div>
  );
}

In this example, the connection state is tracked, but not in the useEffect dependencies. waitingForReconnect is, though. And it's set when the connection is closed, and unset a time later, to trigger a reconnection attempt.

The cleanup triggers a close, as well, so we need to differentiate in the onClose, which we do by seeing if the client has been dereferenced.

As you can see, this approach is rather complex, and it ties the WS lifecycle to the component lifecycle (which is technically ok, if you are doing it at the app level).

However, one major caveat is that it's really easy to run into issues with stale closures. For example, the addMessage has access to the local variable messages, but since addMessage is not passed in as a dependency, you can't call it twice per run of the effect, or it will overwrite the last message. (It's not overwriting, per se; it's actually just overwriting the state with the old, "stale" value of messages, concatenated with the new one. Call it ten times and you'll only see the last value.)

So, you could add addMessage to the dependencies, but then you'd be disconnecting and reconnecting the websocket every render. You could get rid of addMessages, and just move that logic into the effect, but then it would re-run every time you update the messages array (less frequently than on every render, but still too often).

So, coming full circle, I'd recommend setting up your client outside of the app lifecycle. You can use custom hooks to handle incoming messages, or just handle them directly in effects.

Here's an example of that:

import React, { useRef, useState, useEffect } from 'react';

const URL = 'ws://localhost:8888';


function reconnectingSocket(url) {
  let client;
  let isConnected = false;
  let reconnectOnClose = true;
  let messageListeners = [];
  let stateChangeListeners = [];

  function on(fn) {
    messageListeners.push(fn);
  }

  function off(fn) {
    messageListeners = messageListeners.filter(l => l !== fn);
  }

  function onStateChange(fn) {
    stateChangeListeners.push(fn);
    return () => {
      stateChangeListeners = stateChangeListeners.filter(l => l !== fn);
    };
  }

  function start() {
    client = new WebSocket(URL);

    client.onopen = () => {
      isConnected = true;
      stateChangeListeners.forEach(fn => fn(true));
    }

    const close = client.close;

    // Close without reconnecting;
    client.close = () => {
      reconnectOnClose = false;
      close.call(client);
    }

    client.onmessage = (event) => {
      messageListeners.forEach(fn => fn(event.data));
    }

    client.onerror = (e) => console.error(e);

    client.onclose = () => {

      isConnected = false;
      stateChangeListeners.forEach(fn => fn(false));

      if (!reconnectOnClose) {
        console.log('ws closed by app');
        return;
      }

      console.log('ws closed by server');

      setTimeout(start, 3000);
    }
  }

  start();

  return {
    on,
    off,
    onStateChange,
    close: () => client.close(),
    getClient: () => client,
    isConnected: () => isConnected,
  };
}


const client = reconnectingSocket(URL);


function useMessages() {
  const [messages, setMessages] = useState([]);

  useEffect(() => {
    function handleMessage(message) {
      setMessages([...messages, message]);
    }
    client.on(handleMessage);
    return () => client.off(handleMessage);
  }, [messages, setMessages]);

  return messages;
}



export default function App() {

  const [message, setMessage] = useState('');
  const messages = useMessages();
  const [isConnected, setIsConnected] = useState(client.isConnected());

  useEffect(() => {
    return client.onStateChange(setIsConnected);
  }, [setIsConnected]);

  useEffect(() => {
    if (isConnected) {
      client.getClient().send('hi');
    }
  }, [isConnected]);

  function sendMessage(e) {
    e.preventDefault();
    client.getClient().send(message);
    setMessage('');
  }

  return (
    <div>
      <h1>Websocket {isConnected ? 'Connected' : 'Disconnected'}</h1>

      <form onSubmit={sendMessage}>
        <input value={message} onChange={e => setMessage(e.target.value)} />
        <button type="submit">Send</button>
      </form>

      {messages.map(m => <p>{JSON.stringify(m, null, 2)}</p>)}
    </div>
  );
}

Related Query

More Query from same tag