clash-verge-rev/src/components/layout/layout-traffic.tsx

import { useRef } from "react";
import { Box, Typography } from "@mui/material";
import {
  ArrowDownward,
  ArrowUpward,
  MemoryOutlined,
} from "@mui/icons-material";
import { useClashInfo } from "@/hooks/use-clash";
import { useVerge } from "@/hooks/use-verge";
import { TrafficGraph, type TrafficRef } from "./traffic-graph";
import { useLogData } from "@/hooks/use-log-data";
import { useVisibility } from "@/hooks/use-visibility";
import parseTraffic from "@/utils/parse-traffic";
import useSWRSubscription from "swr/subscription";
import { createSockette } from "@/utils/websocket";

interface MemoryUsage {
  inuse: number;
  oslimit?: number;
}

// setup the traffic
export const LayoutTraffic = () => {
  const { clashInfo } = useClashInfo();
  const { verge } = useVerge();

  // whether hide traffic graph
  const trafficGraph = verge?.traffic_graph ?? true;

  const trafficRef = useRef<TrafficRef>(null);
  const pageVisible = useVisibility();

  // https://swr.vercel.app/docs/subscription#deduplication
  // useSWRSubscription auto deduplicates to one subscription per key per entire app
  // So we can simply invoke it here acting as preconnect
  useLogData();

  const { data: traffic = { up: 0, down: 0 } } = useSWRSubscription<
    ITrafficItem,
    any,
    "getRealtimeTraffic" | null
  >(
    clashInfo && pageVisible ? "getRealtimeTraffic" : null,
    (_key, { next }) => {
      const { server = "", secret = "" } = clashInfo!;

      const s = createSockette(
        `ws://${server}/traffic?token=${encodeURIComponent(secret)}`,
        {
          onmessage(event) {
            const data = JSON.parse(event.data) as ITrafficItem;
            trafficRef.current?.appendData(data);
            next(null, data);
          },
          onerror(event) {
            this.close();
            next(event, { up: 0, down: 0 });
          },
        }
      );

      return () => {
        s.close();
      };
    },
    {
      fallbackData: { up: 0, down: 0 },
      keepPreviousData: true,
    }
  );

  /* --------- meta memory information --------- */
  const isMetaCore = verge?.clash_core?.includes("clash-meta");
  const displayMemory = isMetaCore && (verge?.enable_memory_usage ?? true);

  const { data: memory = { inuse: 0 } } = useSWRSubscription<
    MemoryUsage,
    any,
    "getRealtimeMemory" | null
  >(
    clashInfo && pageVisible && displayMemory ? "getRealtimeMemory" : null,
    (_key, { next }) => {
      const { server = "", secret = "" } = clashInfo!;

      const s = createSockette(
        `ws://${server}/memory?token=${encodeURIComponent(secret)}`,
        {
          onmessage(event) {
            const data = JSON.parse(event.data) as MemoryUsage;
            next(null, data);
          },
          onerror(event) {
            this.close();
            next(event, { inuse: 0 });
          },
        }
      );

      return () => {
        s.close();
      };
    },
    {
      fallbackData: { inuse: 0 },
      keepPreviousData: true,
    }
  );

  const [up, upUnit] = parseTraffic(traffic.up);
  const [down, downUnit] = parseTraffic(traffic.down);
  const [inuse, inuseUnit] = parseTraffic(memory.inuse);

  const iconStyle: any = {
    sx: { mr: "8px", fontSize: 16 },
  };
  const valStyle: any = {
    component: "span",
    // color: "primary",
    textAlign: "center",
    sx: { flex: "1 1 56px", userSelect: "none" },
  };
  const unitStyle: any = {
    component: "span",
    color: "grey.500",
    fontSize: "12px",
    textAlign: "right",
    sx: { flex: "0 1 27px", userSelect: "none" },
  };

  return (
    <Box position="relative" onClick={trafficRef.current?.toggleStyle}>
      {trafficGraph && pageVisible && (
        <div style={{ width: "100%", height: 60, marginBottom: 6 }}>
          <TrafficGraph ref={trafficRef} />
        </div>
      )}

      <Box display="flex" flexDirection="column" gap={0.75}>
        <Box display="flex" alignItems="center" whiteSpace="nowrap">
          <ArrowUpward
            {...iconStyle}
            color={+up > 0 ? "secondary" : "disabled"}
          />
          <Typography {...valStyle} color="secondary">
            {up}
          </Typography>
          <Typography {...unitStyle}>{upUnit}/s</Typography>
        </Box>

        <Box display="flex" alignItems="center" whiteSpace="nowrap">
          <ArrowDownward
            {...iconStyle}
            color={+down > 0 ? "primary" : "disabled"}
          />
          <Typography {...valStyle} color="primary">
            {down}
          </Typography>
          <Typography {...unitStyle}>{downUnit}/s</Typography>
        </Box>

        {displayMemory && (
          <Box
            display="flex"
            alignItems="center"
            whiteSpace="nowrap"
            title="Memory Usage"
          >
            <MemoryOutlined {...iconStyle} color="disabled" />
            <Typography {...valStyle}>{inuse}</Typography>
            <Typography {...unitStyle}>{inuseUnit}</Typography>
          </Box>
        )}
      </Box>
    </Box>
  );
};