Building a useSocket Hook in React

Building a useSocket Hook in React

09/09/2021

Ever since the first time I saw someone demo an application using WebSockets, I was completely enthralled. Up to that point, I had always thought of data as something you requested and had returned to you. In this new paradigm shift, I could listen for events from the server and respond to them on the client.

It has actually been a while since I needed to use WebSockets in a live application. I've been using AWS AppSync for most of my APIs recently. AppSync has built-in subscriptions that abstract away having to wire up the WebSockets themselves. However, I recently began working on a project that will ultimately need to run on a local network to combat connectivity issues and AppSync was not an option. I had read about using API Gateway to serve a WebSocket connection. Now, I finally had a chance to try it out.

I will provide an example of how I got the server-side running using the Serverless Stack, API Gateway, DynamoDB and a handful of Lambdas in the future, but today I want to explore how I set up the client-side with a useSocket React hook to handle the WebSocket connection.

One limitation I soon found was that the WebSocket implementation in API Gateway does not play nice with Socket.io, which has been my go-to client for years now. Since the native WebSocket API isn't terribly complicated, I decided to just spin my own WebSocket implementation. There are other libraries out there I could have tried like Sock, but I felt like in this instance I'd probably waste more time researching a library than I would if I just used the native API.

Allow me walk you through how I set up the React hook...

The Provider

We only want to maintain a single WebSocket connection that is shared globally by the application. The React Context API is a perfect place to maintain the WebSocket state. Here, I'm instantiating a new WebSocket and wrapping it in createContext. Finally, I made a simple component to wrap the Context Provider in.

import { useEffect, useState, createContext, ReactChild } from "react";

const ws = new WebSocket("MY_SOCKET_URL");

export const SocketContext = createContext(ws);

interface ISocketProvider {
  children: ReactChild;
}

export const SocketProvider = (props: ISocketProvider) => (
  <SocketContext.Provider value={ws}>{props.children}</SocketContext.Provider>
);

Now the SocketProvider can be used just once near the root of the App like this:

const MyApp = () => {
  return <SocketProvider>{/* Children... */}</SocketProvider>;
};

The Hook

Now that we've got the context and the provider, we can use them inside a custom hook. We'll just pull in the current WebSocket and return it.

import { SocketContext } from "./SocketProvider";
import { useContext } from "react";

export const useSocket = () => {
  const socket = useContext(SocketContext);

  return socket;
};

Emitting Events

Now that we've exposed the socket through our useSocket hook, we just need to invoke it in our component and call the send method on it. Note that I'm wrapping the argument in a JSON.stringify for the WebSocket to transport the message to the server as a string.

import { useSocket } from "./useSocket";

const MyComponent = () => {
  const socket = useSocket();

  return (
    <button
      onClick={() => {
        socket.send(
          JSON.stringify({
            hello: "World",
          })
        );
      }}
    ></button>
  );
};

Listening for Events

Listening for events is a little more complex than emitting them because of their async nature. We essentially need to add an event listener on any messages that come from the WebSocket. We need to filter out anything off-topic and react to any message we care about. This is a good use case for the useCallback hook because we want to have a function we can setup and tear down as we mount and un-mount the component. Using the useCallback hook, we ensure that the function is only instantiated if any properties that are passed into the array, which is the final argument, are changed.

import { useCallback, useEffect } from "react";
import { useSocket } from "./useSocket";

const SomeOtherComponent = () => {
  const socket = useSocket();

  const onMessage = useCallback((message) => {
    const data = JSON.parse(message?.data);
    // ... Do something with the data
  }, []);

  useEffect(() => {
    socket.addEventListener("message", onMessage);

    return () => {
      socket.removeEventListener("message", onMessage);
    };
  }, [socket, onMessage]);

  return (
    <button
      onClick={() => {
        socket.send(
          JSON.stringify({
            hello: "World",
          })
        );
      }}
    ></button>
  );
};

Handling Disconnections

The thing about WebSockets is that they can drop off at any time. Since there is no guarantee that a connection will last indefinitely, we need a way to handle reconnecting to the WebSocket if and when the connection drops. Let's go back to our SocketProvider and wrap the WebSocket instance in a useState. That way, when a socket disconnects, we can listen for the event and create a new WebSocket connection to replace it.

import { useEffect, useState, createContext, ReactChild } from "react";

const webSocket = new WebSocket("MY_SOCKET_URL");

export const SocketContext = createContext(webSocket);

interface ISocketProvider {
  children: React.ReactChild;
}

export const SocketProvider = (props: ISocketProvider) => {
  const [ws, setWs] = useState < WebSocket > webSocket;

  useEffect(() => {
    const onClose = () => {
      setTimeout(() => {
        setWs(new WebSocket(SOCKET_URL));
      }, SOCKET_RECONNECTION_TIMEOUT);
    };

    ws.addEventListener("close", onClose);

    return () => {
      ws.removeEventListener("close", onClose);
    };
  }, [ws, setWs]);

  return (
    <SocketContext.Provider value={ws}>{props.children}</SocketContext.Provider>
  );
};

In Closing

I think this is a really good example of the correct way to thinly wrap an existing API with a React hook.

I opted to use another component on top of useSocket that abstracted away the callback and event emitting so that I wasn't parsing JSON inside components or using useEffects to manage adding and removing event listeners. That is totally a matter of preference, though.

All in all, I'm pretty happy with this solution. How do you handle server events in your React applications?

Never miss an update!

Join our newsletter

Tyson Cadenhead
Tyson Cadenhead leads the dev team at Vestaboard

He has a passion for Functional Programming, GraphQL, Serverless architecture and React. When he's not writing code or working with his team, Tyson enjoys drawing, playing guitar, growing vegetables and spending time with his wife, two boys and his dog and cat.