在 Next.js 中使用 WebSocket

Question

I'm wondering what's the best approach to connect into my WebSockets server using Next.js' pages? I'd like that the user could navigate through the pages with one connection, and when he closes the page he also closes the WebSockets connection. I tried using React's Context API:

const WSContext = createContext(null);

const Wrapper = ({ children }) => {
  const instance = WebSocket("ws://localhost:3000/ws");

  return <WSContext.Provider value={instance}>{children}</WSContext.Provider>;
};

export const useWS = () => useContext(WSContext);

export default Wrapper;

It works great, but it does not when it comes to creating a connection. The basic new WebSocket syntax is not working, so I must use a third-party library, like react-use-websocket which I dislike. What's also disturbing is that I cannot close the connection. The Context simply does not know when the page is closed, and also the library does not provide a hook for closing connections.

I would like to know what's the best approach when it comes to handling WebSockets connection in Next.js.

Answer 1

There are multiple things that need to be done in order for ws to work on Next.js.

Firstly, it is important to realize where do I want my ws code to run. React code on Next.js runs in two environments: On the server (when building the page or when using ssr) and on the client.

Making a ws connection at page build time has little utility, thats why I will cover client-side ws only.

The global Websocket class is a browser only feature and is not present on the server. Thats why we need to prevent any instantiation until the code is run in the browser. One simple way to do so would be:

export const isBrowser = typeof window !== "undefined";

export const wsInstance = isBrowser ? new Websocket(...) : null;

Also, you do not need to use react context for holding the instance, it is perfectly possible to keep it at global scope and import the instance, unless you wish to open the connection lazily.

If you still decide to use react context (or initialize the ws client anywhere in the react tree), it is important to memoize the instance, so that it isn't created on every update of your react node.

const wsInstance = useMemo(() => isBrowser ? new Websocket(...) : null, []);

or

const [wsInstance] = useState(() => isBrowser ? new Websocket(...) : null);

Any event handler registrations created in react should be wrapped inside a useEffect with a return function which removes the event listener, this is to prevent memory leaks, also, a dependency array should be specified. If your component is unmounted, the useEffect hook will remove the event listener as well.

If you wish to reinitailize the ws and dispose of the current connection, then it is possible to do something similar to the following:

const [wsInstance, setWsInstance] = useState(null);

// Call when updating the ws connection
const updateWs = useCallback((url) => {
   if(!browser) return setWsInstance(null);
   
   // Close the old connection
   if(wsInstance?.readyState !== 3)
     wsInstance.close(...);

   // Create a new connection
   const newWs = new WebSocket(url);
   setWsInstance(newWs);
}, [wsInstance])


// (Optional) Open a connection on mount
useEffect(() => {
 if(isBrowser) { 
   const ws = new WebSocket(...);
   setWsInstance(ws);
 }

 return () => {
  // Cleanup on unmount if ws wasn't closed already
  if(ws?.readyState !== 3) 
   ws.close(...)
 }
}, [])

Answer 2

I configured context like this:

import { createContext, useContext, useMemo, useEffect, ReactNode } from 'react';

type WSProviderProps = { children: ReactNode; url: string };

const WSStateContext = createContext<WebSocket | null>(null);

function WSProvider({ children, url }: WSProviderProps): JSX.Element {
  const wsInstance = useMemo(
    () => (typeof window != 'undefined' ? new WebSocket(`ws://127.0.0.1:8001/ws${url}`) : null),
    []
  );

  useEffect(() => {
    return () => {
      wsInstance?.close();
    };
  }, []);

  return <WSStateContext.Provider value={wsInstance}>{children}</WSStateContext.Provider>;
}

function useWS(): WebSocket {
  const context = useContext(WSStateContext);

  if (context == undefined) {
    throw new Error('useWS must be used within a WSProvider');
  }

  return context;
}

export { WSProvider, useWS };