Managing Websockets in React

Published: Wednesday, March 10, 2021

I recently worked on a react project that needed to listen for events from a backend service, and perform updates to state. Websockets are the obvious choice for this. There doesn't seem to be a lot of good, straightforward examples of using websockets in modern (hooks-based) react. In this short article, I will be sharing a simple implementation I have been using to respond to events on a websocket.

Overview

We will create a new React Context (WebsocketContext), to contain logic to manage websocket connections, and listen for incoming messages. Upon receiving a message, the WebsocketContext will broadcast that message event to every subscriber (listener). Components (or other contexts) can subscribe (listen) to messages from the WebsocketContext, by adding a listener function to a collection of listeners stored within our WebsocketContext.

Create Websocket Context

The first step was to leverage the React Context API and create a new context to hold the logic and state for our websocket connections. Here is what the context looks like, including supporting models and types:


export enum WebSocketActionType {
    Unknown = 0,
    NewNotification = 1,
    DirectMessage = 2,
}

export class WebsocketMessage {
    actionType: WebSocketActionType;
    data: any;
}

export interface IWebsocketListener {
    id: number;
    messageHandler: (msg: WebsocketMessage) => void;
}

let listeners: IWebsocketListener[] = [];

export interface IWebSocketContext {
    initializeWebSocket: () => void;
    disconnectWebSocket: () => void;
    addListener: (listener: (msg: WebsocketMessage) => void) => number;
    removeListener: (id: number) => void;
}

const WebsocketContext = React.createContext({} as IWebSocketContext);

const WebsocketProvider: React.FC = ({ children }) => {
    const websocket = useRef(null);

    let autoReconnectInterval = 250; // in ms

    const reconnectToWebSocket = (autoReconnectInterval: number) => {
        const maximumConnectionTimeout = 10000;
        setTimeout(initializeWebSocket, Math.min(maximumConnectionTimeout, autoReconnectInterval));
    };

    const initializeWebSocket = () => {
        websocket.current = new WebSocket(`wss://your-websocket-url`);

        websocket.current.onopen = (event: Event) => {
            autoReconnectInterval = 250; //reset the interval for reconnecting
        };

        websocket.current.onmessage = (message: MessageEvent) => {
            const decodedMessage = JSON.parse(message.data);
            handleMessage(decodedMessage);
        };

        websocket.current.onclose = (event: CloseEvent) => {
            switch (event.code) {
                case 1000: //closed normally
                    console.debug('WebSocket closed normally');
                    break;
                default:
                    autoReconnectInterval += autoReconnectInterval; //increment the interval for reconnecting
                    reconnectToWebSocket(autoReconnectInterval);
                    break;
            }
        };

        websocket.current.onerror = (error: any) => {
            switch (error.code) {
                case 'ECONNREFUSED':
		    //increment the interval for reconnecting
                    autoReconnectInterval = autoReconnectInterval += autoReconnectInterval; 
                    reconnectToWebSocket(autoReconnectInterval);
                    break;
                default:
                    console.error(`WebSocket encountered error: ${error}`);
                    break;
            }
        };
    };

    const handleMessage = (message: WebsocketMessage) => {
        console.debug('[WebSocket-IncomingMessage]', message);
        listeners.forEach((x) => {
            x.messageHandler(message);
        });
    };

    const disconnectWebSocket = () => {
        if (!!websocket.current) websocket.current.close(1000);
    };

    const addListener = (messageHandler: (msg: WebsocketMessage) => void): number => {
        const temp = [...listeners];
        let id: number;
        do {
            id = Math.floor(Math.random() * 10000) + 1;
        } while (temp.map((x) => x.id).includes(id)); // avoid duplicates
        temp.push({ id, messageHandler });
        listeners = temp;
        return id;
    };

    const removeListener = (id: number) => {
        const temp = [...listeners];
        const idToRemove = temp.findIndex((x) => x.id === id);
        temp.splice(idToRemove, 1);
        listeners = temp;
    };

    return (
        <WebsocketContext.Provider
            value={{
                initializeWebSocket: initializeWebSocket,
                disconnectWebSocket: disconnectWebSocket,
                addListener,
                removeListener,
            }}
        >
            {children}
        </WebsocketContext.Provider>
    );
};

const WebsocketConsumer = WebsocketContext.Consumer;
export { WebsocketProvider, WebsocketConsumer, WebsocketContext };
	

A few notes on the above:

  • The collection of listeners is maintained outside of the context/React. React's state management doesn't play nice with websockets evidently.
  • This websocket implementation has automatic reconnect functionality built in, to try and reinstate a connection in the event something bad happens

Initialize Websocket

Before we can use the websockets, we need to call the initializeWebSocket function on our new context. In my case, I called this after a successful login to my application.

Using Websockets

We can use our websocket implementation in both components and contexts. First we need to get references to addListener and removeListener functions defined in our context:

const { addListener, removeListener } = useContext(WebsocketContext);

We will need to store the websocket listener ID in our component, so let's add state for that:

const [websocketId, setWebsocketId] = useState(-1);

Then, we will add a useEffect to act as a constructor/destructor in our component. This will allow us to begin listening for events on our websocket as soon as the component mounts. When the component unmounts, the onComponentDestroy function will be called, as it's being returned from our "constructor" useEffect. Don't forget to remove the listeners before your component unmounts!! (Memory leaks and weird bugs are no fun)

const initWebsocketListener = () => {
	const websocketListenerId = addListener((msg: WebsocketMessage) => {
		if (msg.actionType === WebsocketActionType.NewNotification) reloadNotifications();
		if (msg.actionType === WebsocketActionType.DirectMessage) showMessageNotification(msg);
	};
	setWebsocketId(websocketListenerId);
}

const onComponentDestroy = () => {
	 // remove our listener from the context's list of listeners
	removeListener(websocketListenerId);
	setWebsocketId(-1);
}

useEffect(() => {
	initWebsocketListener();

	return onComponentDestroy;
	// eslint-disable-next-line
}, []);

Room for improvement

This is a very simple implementation, with some room for improvement! Specifically, I would like to eventually find a better (more react-friendly) way to hold the collection of listeners. Furthermore, adding functionality in the future to send websocket messages can be added.

I hope this was helpful! As always, let me know if you have any questions or suggestions!