ontime/apps/client/src/common/utils/socket.ts

import {
  ApiActionTag,
  Log,
  MaybeNumber,
  MessageTag,
  RefetchKey,
  RuntimeStore,
  WsPacketToClient,
  WsPacketToServer,
} from 'ontime-types';

import { isProduction, websocketUrl } from '../../externals';
import {
  APP_SETTINGS,
  CLIENT_LIST,
  CSS_OVERRIDE,
  CURRENT_RUNDOWN_QUERY_KEY,
  CUSTOM_FIELDS,
  PROJECT_DATA,
  REPORT,
  RUNDOWN,
  RUNTIME,
  TRANSLATION,
  URL_PRESETS,
  VIEW_SETTINGS,
  getRundownQueryKey,
  PROJECT_RUNDOWNS,
} from '../api/constants';
import { invalidateAllCaches } from '../api/utils';
import { ontimeQueryClient } from '../queryClient';
import {
  getClientId,
  getClientName,
  setClientId,
  setClientName,
  setClientRedirect,
  setClients,
} from '../stores/clientStore';
import { addDialog } from '../stores/dialogStore';
import { addLog } from '../stores/logger';
import { patchRuntime, patchRuntimeProperty } from '../stores/runtime';

let websocket: WebSocket | null = null;
let reconnectTimeout: NodeJS.Timeout | null = null;
const socketConfig = {
  reconnectBaseInterval: 1000, // 1 second
  reconnectMaxInterval: 30000, // 30 seconds
  reconnectMinInterval: 500, // 0.5 seconds
  reconnectJitter: 0.25,
  offlineAttemptsThreshold: 2, // when we consider the client disconnected
} as const;

export const getConnectionState = () => hasConnected;
export const getReconnectAttempts = () => reconnectAttempts;
let hasConnected = false;
let reconnectAttempts = 0;

export const connectSocket = () => {
  websocket = new WebSocket(websocketUrl);

  const preferredClientName = getClientName();

  websocket.onopen = () => {
    const isReconnect = hasConnected;
    if (reconnectTimeout) {
      clearTimeout(reconnectTimeout);
      reconnectTimeout = null;
    }
    hasConnected = true;
    reconnectAttempts = 0;

    sendSocket(MessageTag.ClientSet, {
      type: 'ontime',
      origin: window.location.origin,
      path: window.location.pathname + window.location.search,
      name: preferredClientName,
    });

    if (isReconnect) {
      invalidateAllCaches();
    }
    setOnlineStatus(true);
  };

  websocket.onclose = () => {
    console.warn('WebSocket disconnected');
    if (reconnectTimeout) {
      clearTimeout(reconnectTimeout);
      reconnectTimeout = null;
    }

    const exponentialDelay = Math.min(
      socketConfig.reconnectBaseInterval * 2 ** reconnectAttempts,
      socketConfig.reconnectMaxInterval,
    );
    const jitterOffset = exponentialDelay * socketConfig.reconnectJitter * (Math.random() * 2 - 1);
    const delay = Math.max(socketConfig.reconnectMinInterval, Math.round(exponentialDelay + jitterOffset));

    reconnectTimeout = setTimeout(() => {
      reconnectTimeout = null;
      if (reconnectAttempts > socketConfig.offlineAttemptsThreshold) {
        setOnlineStatus(false);
      }
      console.warn(`WebSocket: reconnecting now (#${reconnectAttempts + 1}, waited ${delay}ms)`);
      if (websocket && websocket.readyState === WebSocket.CLOSED) {
        reconnectAttempts += 1;
        connectSocket();
      }
    }, delay);
  };

  websocket.onerror = (error) => {
    console.error('WebSocket error:', error);
  };

  websocket.onmessage = async (event) => {
    try {
      const data = JSON.parse(event.data) as WsPacketToClient;

      const { tag, payload } = data;

      if (!tag) {
        return;
      }

      switch (tag) {
        case MessageTag.Pong: {
          const offset = (new Date().getTime() - new Date(payload).getTime()) * 0.5;
          patchRuntimeProperty('ping', offset);
          updateDevTools({ ping: offset });
          break;
        }
        case MessageTag.ClientInit: {
          setClientId(payload.clientId);
          if (!preferredClientName) {
            setClientName(payload.clientName);
          }
          break;
        }

        case MessageTag.ClientRename: {
          const id = getClientId();
          if (payload.target === id) {
            setClientName(payload.name);
          }
          break;
        }

        case MessageTag.ClientRedirect: {
          const id = getClientId();
          if (payload.target === id) {
            setClientRedirect(payload.path);
          }
          break;
        }

        case MessageTag.ClientList: {
          setClients(payload);
          if (!isProduction) {
            ontimeQueryClient.setQueryData(CLIENT_LIST, payload);
          }
          break;
        }

        case MessageTag.Dialog: {
          if (payload.dialog === 'welcome') {
            addDialog('welcome');
          }
          break;
        }

        case MessageTag.Log: {
          addLog(payload as Log);
          break;
        }
        case MessageTag.RuntimeData: {
          patchRuntime(payload);
          updateDevTools(payload);
          break;
        }
        case MessageTag.Refetch: {
          // the refetch message signals that the rundown has changed in the server side
          const { target, revision, rundownId } = payload;
          switch (target) {
            case RefetchKey.All:
              invalidateAllCaches();
              break;
            case RefetchKey.CustomFields:
              ontimeQueryClient.invalidateQueries({ queryKey: CUSTOM_FIELDS });
              break;
            case RefetchKey.ProjectData:
              ontimeQueryClient.invalidateQueries({ queryKey: PROJECT_DATA });
              break;
            case RefetchKey.Report:
              ontimeQueryClient.invalidateQueries({ queryKey: REPORT });
              break;
            case RefetchKey.Rundown: {
              maybeInvalidateRundownCache(revision, rundownId);
              break;
            }
            case RefetchKey.UrlPresets:
              ontimeQueryClient.invalidateQueries({ queryKey: URL_PRESETS });
              break;
            case RefetchKey.ViewSettings:
              ontimeQueryClient.invalidateQueries({ queryKey: VIEW_SETTINGS });
              break;
            case RefetchKey.CssOverride:
              ontimeQueryClient.invalidateQueries({ queryKey: CSS_OVERRIDE });
              break;
            case RefetchKey.Translation:
              ontimeQueryClient.invalidateQueries({ queryKey: TRANSLATION });
              break;
            case RefetchKey.Settings:
              ontimeQueryClient.invalidateQueries({ queryKey: APP_SETTINGS });
              break;
            case RefetchKey.ProjectRundowns:
              ontimeQueryClient.invalidateQueries({ queryKey: PROJECT_RUNDOWNS });
              break;
            default: {
              target satisfies never;
              break;
            }
          }
          break;
        }
        default: {
          tag satisfies never;
          break;
        }
      }
    } catch (_) {
      // ignore unhandled
    }
  };
};

export function maybeInvalidateRundownCache(revision: MaybeNumber, rundownId?: string) {
  if (!rundownId) {
    // we omit rundownId to signify invalidate all rundowns
    ontimeQueryClient.invalidateQueries({ queryKey: RUNDOWN });
    ontimeQueryClient.invalidateQueries({ queryKey: CURRENT_RUNDOWN_QUERY_KEY, exact: true });
    return;
  }

  // skip if we dont recognise the ID the revision is lower
  const queryKey = getRundownQueryKey(rundownId);
  const cachedRundown = ontimeQueryClient.getQueryData<{ revision: number }>(queryKey);

  if (revision === cachedRundown?.revision) {
    // we already have the latest change
    return;
  }

  ontimeQueryClient.invalidateQueries({ queryKey, exact: true });

  // keep current alias in sync with the ID-based cache
  const loadedRundownId = ontimeQueryClient.getQueryData<{ loaded: string }>(PROJECT_RUNDOWNS)?.loaded;
  if (!loadedRundownId || loadedRundownId === rundownId) {
    ontimeQueryClient.invalidateQueries({ queryKey: CURRENT_RUNDOWN_QUERY_KEY, exact: true });
  }
}

export function sendSocket<T extends MessageTag | ApiActionTag>(
  tag: T,
  payload: T extends MessageTag ? Pick<WsPacketToServer & { tag: T }, 'payload'>['payload'] : unknown,
): void {
  if (websocket && websocket.readyState === WebSocket.OPEN) {
    websocket.send(JSON.stringify({ tag, payload }));
  }
}

function updateDevTools(newData: Partial<RuntimeStore>) {
  if (!isProduction) {
    ontimeQueryClient.setQueryData(RUNTIME, (oldData: RuntimeStore) => ({
      ...oldData,
      ...newData,
    }));
  }
}

/**
 * Allows setting the status of the client
 * We leverage the ping as an indication of the client's online status
 * @example ping < 0 - client is offline
 * @example ping > 0 -> client is online
 */
function setOnlineStatus(status: boolean) {
  const derivedPing = status ? 1 : -1;
  patchRuntimeProperty('ping', derivedPing);
  updateDevTools({ ping: derivedPing });
}