Websockets with React & Express [Part-2]

In Part-1 we setup a websocket server which can send individual/broadcast messages to client and it runs alongside the express server. Part-2 will focus on the client side code.

To have the websocket connection code as reusable and to use it easily in our components it would be good to write it as a custom react hook.

Our hook should do the following to start with

• connect to socket and maintain the connection state
• send and receive messages
• handle parsing and stringify-ing of messages
• have a retry logic in case the connection gets terminated

// webSocketHook.js
import { useState, useEffect } from "react";

// define a custom hook
// accept the url to connect to
// number of times the hook should retry a connection
// the interval between retries
function useWebSocketLite({
  socketUrl,
  retry: defaultRetry = 3,
  retryInterval = 1500
}) {
  // message and timestamp
  const [data, setData] = useState();
  // send function
  const [send, setSend] = useState(() => () => undefined);
  // state of our connection
  const [retry, setRetry] = useState(defaultRetry);
  // retry counter
  const [readyState, setReadyState] = useState(false);

  useEffect(() => {
    const ws = new WebSocket(socketUrl);
    ws.onopen = () => {
      console.log('Connected to socket');
      setReadyState(true);

      // function to send messages
      setSend(() => {
        return (data) => {
          try {
            const d = JSON.stringify(data);
            ws.send(d);
            return true;
          } catch (err) {
            return false;
          }
        };
      });

      // receive messages
      ws.onmessage = (event) => {
        const msg = formatMessage(event.data);
        setData({ message: msg, timestamp: getTimestamp() });
      };
    };

    // on close we should update connection state
    // and retry connection
    ws.onclose = () => {
      setReadyState(false);
      // retry logic
      if (retry > 0) {
        setTimeout(() => {
          setRetry((retry) => retry - 1);
        }, retryInterval);
      }
    };
     // terminate connection on unmount
    return () => {
      ws.close();
    };
  // retry dependency here triggers the connection attempt
  }, [retry]); 

  return { send, data, readyState };
}

// small utilities that we need
// handle json messages
function formatMessage = (data) => {
  try {
    const parsed = JSON.parse(data);
    return parsed;
  } catch (err) {
    return data;
  }
};

// get epoch timestamp
function getTimestamp() {
  return new Date().getTime();
}

export default useWebSocketLite;

Time to use it in a component, Let's have the following things in the component

• list of all messages, send and received
• an indicator for message direction
• a form to post a message to our sever

import React, { useState, useEffect, useRef } from 'react';
import useWebSocketLite from './webSocketHook.js;'

// prettify
const sendTag = (message) => <span>&#11014;: {message}</span>;
const receiveTag = (message) => <span>&#11015;: {message}</span>;

function App() {
  const [messagesList, setMessagesList] = useState([
    <span>Messages will be displayed here</span>
  ]);
  const txtRef = useRef();

  // use our hook
  const ws = useWebSocket({
    socketUrl: 'ws://localhost:3000'
  });

  // receive messages
  useEffect(() => {
    if (ws.data) {
      const { message } = ws.data;
      setMessagesList((messagesList) =>
        [].concat(receiveTag(message), messagesList)
      );
    }
  }, [ws.data]);

  // send messages
  const sendData = () => {
    const message = txtRef.current.value || '';
    if (message) {
      setMessagesList((messagesList) =>
        [].concat(sendTag(message), messagesList)
      );
      ws.send(message);
    }
  };

  // a simple form
  return (
    <div>
       <div>Connection State: {ws.readyState ? 'Open' : 'Closed'}</div>

      <div>      
       <form>
          <label>Message (string or json)</label>
          <textarea name='message' rows={4} ref={txtRef} />
          <input type='button' onClick={sendData} value='Send' />
        </form>
      </div>

      <div style={{ maxHeight: 300, overflowY: 'scroll' }}>
        {messagesList.map((Tag, i) => (
          <div key={i}>{Tag}</div>
        ))}
      </div>

    </div>
  );
}

export default App;

Mount this component and watch for the ping-pong/broadcast messages. I added milligram and some styling to make it pretty.

In Part-3 we'll add some more functionality to both client and server.