All files / src/hooks/notifications useNotificationConnection.ts

100% Statements 25/25
100% Branches 8/8
100% Functions 6/6
100% Lines 24/24

Press n or j to go to the next uncovered block, b, p or k for the previous block.

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81                    66x 33x 33x 33x               27619x       27619x 120x 120x 120x       120x 120x     120x     27619x   3313x   2533x     2533x       1656x   877x     27619x 27619x 299x                 120x                   27619x   27619x    
import { useCallback, useMemo } from "react";
import useWebSocket, { Options } from "react-use-websocket";
import { useProjectId } from "@/hooks/project.ts";
import { useEventCallback } from "usehooks-ts";
 
export type NotificationMessage = {
  message_type: NotificationMessageType;
  data: never;
};
 
export enum NotificationMessageType {
  TASK_COMPLETED = "task/completed",
  TASK_UPDATED = "task/updated",
  TASK_CANCELLED = "task/cancelled"
}
 
export function useNotificationConnection(
  onMessage: (message: never) => void,
  subscribedMessageType?: NotificationMessageType,
) {
  // @ts-expect-error Vite environment variables are dynamically injected at build time
  const socketUrl = import.meta.env.VITE_NOTIFICATION_ENDPOINT_URL;
 
  // Calculates the period (in milliseconds) to wait before connecting again, using exponential backoff with some
  // random jitter to prevent all clients from reconnecting at the same time.
  const retryHandler = useCallback((retryCount: number) => {
    const baseWaitPeriod = 250; // 0.25 seconds
    const maxWaitPeriod = 10000; // 10 seconds
    const exponentialWaitPeriod = Math.min(
      Math.pow(2, retryCount) * baseWaitPeriod,
      maxWaitPeriod,
    );
    const jitter = Math.random() * 0.1 * exponentialWaitPeriod; // Random jitter of up to 10% of the wait period
    const sign = Math.random() < 0.5 ? -1 : 1; // Randomly choose the sign of the jitter
 
    // Wait = exponentialWaitPeriod +/- 10%
    return exponentialWaitPeriod + sign * jitter;
  }, []);
 
  const messageHandler = useEventCallback((message: MessageEvent) => {
    // Ping/pong messages are handled by the library
    if (message.data === "pong") return;
 
    const notificationMessage = JSON.parse(message.data) as NotificationMessage;
 
    // Only deliver messages that have been subscribed to
    if (
      subscribedMessageType &&
      notificationMessage.message_type !== subscribedMessageType
    )
      return;
 
    onMessage(notificationMessage.data);
  });
 
  const projectId = useProjectId();
  const options: Options = useMemo(
    () => ({
      heartbeat: {
        message: "ping",
        returnMessage: "pong",
        timeout: 10000, // 10 seconds
        interval: 5000, // 5 seconds
      },
      onMessage: messageHandler,
      share: true,
      shouldReconnect: () => true, // Always reconnect
      reconnectInterval: retryHandler,
      retryOnError: true,
      queryParams: {
        flowsheetId: projectId,
      },
    }),
    [messageHandler, projectId, retryHandler],
  );
 
  const { readyState } = useWebSocket(socketUrl, options);
 
  return { readyState };
}