useWebsocket.ts

useWebsocket.ts

import * as React from "react";

function preventReconnect(timeoutIDRef: React.MutableRefObject<number>) {
  if (timeoutIDRef.current !== -1) {
    window.clearInterval(timeoutIDRef.current);
  }
}

function isEventSourceActive(eventSource: WebSocket) {
  return eventSource && eventSource.readyState === WebSocket.OPEN;
}

function closeActiveEventSource(eventSource: WebSocket) {
  if (isEventSourceActive(eventSource)) {
    eventSource.close();
  }
}
// TODO needs testing
export default function useWebsocket(
  url: string,
  onMessage: (message: any) => void,
  autoRestartTimeout = 15
) {
  const [eventSource, setEventSource] = React.useState<WebSocket>(null);
  const timeoutIDRef = React.useRef<number>(-1);
  const [WSState, setWSState] = React.useState<number>(WebSocket.CONNECTING);

  const connectToWebSocket = React.useCallback(
    function connectToWebSocket(wsurl: string) {
      closeActiveEventSource(eventSource);
      const newEventSource = new WebSocket(wsurl);
      setWSState(WebSocket.CONNECTING);
      newEventSource.onopen = () => {
        setWSState(WebSocket.OPEN);
      };
      setEventSource(newEventSource);
    },
    [eventSource]
  );

  const reconnect = React.useCallback(
    function reconnect() {
      setWSState(WebSocket.CLOSED);

      preventReconnect(timeoutIDRef);

      if (autoRestartTimeout === 0) {
        return;
      }

      const timeout = autoRestartTimeout || 15;

      timeoutIDRef.current = window.setTimeout(
        connectToWebSocket,
        timeout * 1000
      );
    },
    [autoRestartTimeout, connectToWebSocket]
  );

  React.useEffect(() => {
    if (!eventSource) {
      return;
    }
    eventSource.addEventListener("message", onMessage);
    return () => {
      eventSource.removeEventListener("message", onMessage);
    };
  }, [eventSource, onMessage]);

  React.useEffect(() => {
    if (!eventSource) {
      return;
    }

    const errorHandler = () => {
      if (eventSource.readyState === WebSocket.CLOSING) {
        setWSState(WebSocket.CLOSING);
      } else {
        reconnect();
      }
    };

    eventSource.addEventListener("close", reconnect);
    eventSource.addEventListener("error", errorHandler);

    return () => {
      eventSource.removeEventListener("close", reconnect);
      eventSource.removeEventListener("error", errorHandler);
    };
  }, [eventSource, reconnect]);

  React.useEffect(
    function connect() {
      connectToWebSocket(url);
      return () => {
        preventReconnect(timeoutIDRef);
        closeActiveEventSource(eventSource);
      };
    },
    // Fine as it is with [url]
    [url]
  );

  return { WSState, eventSource };
}