Using WebSockets with React.js, the right way (no library needed)
Tags: javascript, react, webdev, beginners
Published: 2021-07-29 10:17:08

### TL;DR
In this post I introduce useful custom React.js hooks that take websocket clients to the next level.

### Introduction
In the project I'm currently working on, I have a React.js frontend and a WebSocket server that need to be connected. I spent weeks trying to figure out the best way to use websockets, and I wanted the share the things I learned here.

The code solutions I introduce:
1. Why using the `useReducer()` hook when working with WebSockets?
2. My custom `useSession()` hook
3. My usage of the `useCallback()` hook.
4. Auto-reconnect features with the custom `useDisconnectionHandler()` hook.
Bonus: Auto-reconnect on page refresh when needed.

## The useReducer hook

When I first tried to implement my state management system and update it properly when a message was received, it was a disaster.

My `GameContextProvider` component, responsible for handling such events, looked like this:
```js
// GameContextProvider.js

const GameContextProvider = ({ children }) => {
  const [isStarted, setStarted] = useState(false);
  const [isConnected, setConnected] = useState(false);
  const [isJudge, setIsJudge] = useState(false);
  const [judge, setJudge] = useState('');
  const [question, setQuestion] = useState('');
  const [deck, setDeck] = useState([]);
  const [showEndScreen, setEndScreenShown] = useState(false);
  const [scoreboard, setScoreboard] = useState([]);
  ........ 
  // Much more state!
  .....
}

```
Then, when I wanted to handle websocket messages, the handler looked like this:

```js
// GameContextProvider.js

const onMessage = (ev) => {
  const data = JSON.parse(ev.data); 
  if (data.question) { // a round is started
    setJudge(data.judge);
    setIsJudge(data.isJudge);
    setQuestion(data.question);
  }
   ...... // super long, unreadable message handler
}
```
### The Solution
I attached a 'context' string to each of my messages in the server, and used this string to dispatch an action in the useReducer hook.
For example, I had a 'JOINED' context, 'GAME_STARTED', 'ROUND_STARTED', 'GAME_ENDED', etc...

then, my `GameContextProvider` looked like this:
```js
// GameContextProvider.js
const [state, dispatch] = useReducer(reducer, initialState);

const onMessage = (ev) => {
  const data = JSON.parse(ev.data); 
  if (data.context) 
    dispatch({ type: data.context, payload: data })
}
```
Simple and clean!

In addition, this follows the **single responsibility** rule. Now the component's responsibility was to wire the state and the websocket technology for the rest of the application to use.

## The useSession hook

Before I splitted the WebSocket utilities to a custom hook, my context provider had a messy, unreadable code that took care of the websocket events.

```js
// GameContextProvider.js
const [ws, setWebsocket] = useState(null)
const join = (gameCode, name) => {
  const URL = `${process.env.REACT_APP_WS_URL}?code=${gameCode}&name=${name}`
  setWebsocket(() => {
    const ws = new WebSocket(URL);
    ws.onmessage = onMessage;
    ws.onclose = () => {
      dispatch({ type: 'DISCONNECTED' })
    };
    return ws;
  })
}
```

On the surface, this approach looks OK.
but what if I wanted to check the game state on disconnection? If I was to register the function as is, when the value of the state updates, the function would not update!

### The Solution
I created a custom hook that handled the websocket utilities. (Note - by that time I refactored my project to TypeScript)
```ts
// websocketUtils.ts

export const useSession = (
  onOpen: OpenHandler, 
  onMessage: MessageHandler, 
  onClose: CloseHandler
): SessionHook => {
  const [session, setSession] = useState(null as unkown as Websocket);
  const updateOpenHandler = () => {
    if (!session) return;
    session.addEventListener('open', onOpen);
    return () => {
      session.removeEventListener('open', onOpen);
    };
  };

  const updateMessageHandler = () => {
    if (!session) return;
    session.addEventListener('message', onMessage);
    return () => {
      session.removeEventListener('message', onMessage);
    };
  };

  const updateCloseHandler = () => {
    if (!session) return;
    session.addEventListener('close', onClose);
    return () => {
      session.removeEventListener('close', onClose);
    };
  };

  useEffect(updateOpenHandler, [session, onOpen]);
  useEffect(updateMessageHandler, [session, onMessage]);
  useEffect(updateCloseHandler, [session, onClose]);

   .... // connect, sendMessage utils
}

```

This was great! But for some reason, the website's performance was decreased dramatically.

## The useCallback hook

To be honest, I had no idea how this hook worked until last week, when I finally figured out the solution.
As it turns out, my open, message, and close handlers were updated on every re-render of the app (!), meaning a few times per second.

When I debugged the application, I tried to test out the affect of the `useCallback` hook at my performance. as it turned out, the callback hook was only updating the function when one of its dependencies changed, meaning once in minutes!

This improved the performance of my application dramatically.

```ts
// GameContextProvider.tsx
const disconnectHandler = useCallback(() => {
  if (state.gameStatus !== GameLifecycle.STOPPED) // unexpected disconnection!
    console.log('unexpected disconnection')
}, [state.gameStatus])
```

## My Custom disconnection handler hook

In the current version of my project, I wanted to develop a feature - on unexpected disconnection, try to reconnect!

I made the changes to my API and was ready to implement them in my React.js client.

As it turned out, this is possible:
```ts
// eventHandlers.ts
export const useConnectionPauseHandler(
  state: IGameData,
  dispatch: React.Dispatch<any>
) => {
  const [connectFn, setConnectFn] = useState<ConnectFN>(
    null as unknown as ConnectFN
  );

  const disconnectCallback = useCallback(() => {
    if (state.connectionStatus !== ConnectionLifecycle.RESUMED)
      dispatch({ type: 'DISCONNECTED' });
  }, [dispatch, state.connectionStatus]);

  const pauseCallback = useCallback(() => {
    if (...) {
      // disconnection is expected, or an error is prevting the connection from reconnecting
      console.log('expected disconnection');
      dispatch({ type: 'DISCONNECTED' });
    } else if (...) {
      // connection is unexpected, and not attempting reconnection
      console.log('unexpected disconnection');
      dispatch('SESSION_PAUSED');
      if (connectFn) connectFn(state.gameCode!, null, state.playerId);
      setTimeout(disconnectCallback, 30 * 1000);
    }
  }, [
    disconnectCallback,
    dispatch,
    connectFn,
    state.gameCode,
    state.playerId,
    state.connectionStatus,
    state.gameStatus,
  ]);

  const registerConnectFunction = useCallback((fn: ConnectFN) => {
    setConnectFn(() => fn); // do this to avoid confusing the react dispatch function
  }, []);

  return [registerConnectFunction, pauseCallback];
}

// GameContextProvider.tsx
  const [setConnectFn, onClose] = useConnectionPauseHandler(state, dispatch);
  const [connect, sendMessage] = useSession(
    onOpen,
    onMessage,
    onClose
  );

  useEffect(() => {
    console.log('wiring everything...');
    setConnectFn(connect);
  }, [setConnectFn, connect]);

```

The feature worked like magic.

## Bonus

This is a component that saved the connection credentials if the page is refreshed. Can you figure out a way to refactor it to hooks?

```tsx
export default class LocalStorageConnectionRestorer extends Component<Wrapper> {
  static contextType = GameContext;
  state = { isReady: false };
  saveValuesBeforeUnload = () => {
    const { connectionStatus, showEndScreen, gameCode, playerId, close } =
      this.context;
    if (connectionStatus === ConnectionLifecycle.RESUMED && !showEndScreen) {
      // going away before game is over
      console.log('saving reconnection before unmount', gameCode, playerId);
      LocalStorageUtils.setValues(gameCode!, playerId!);
      close();
    }
  };
  componentDidMount() {
    const [gameCode, playerId] = LocalStorageUtils.getValues();
    if (gameCode && playerId) {
      console.log('attempting reconnection after render');
      this.context.reconnect(gameCode, playerId);
      LocalStorageUtils.deleteValues();
    }
    this.setState({ isReady: true });
    window.addEventListener('beforeunload', this.saveValuesBeforeUnload);
  }
  componentWillUnmount() {
    window.removeEventListener('beforeunload', this.saveValuesBeforeUnload);
  }
  render() {
    return this.state.isReady ? (
      this.props.children
    ) : (
      <div className="flex items-center justify-center">Loading...</div>
    );
  }
}
```

### View The Full Source Code
{% github itays123/partydeck no-readme %}