105 lines
2.7 KiB
JavaScript
105 lines
2.7 KiB
JavaScript
import { useEffect, useRef } from 'react';
|
|
|
|
function buildWsUrl() {
|
|
if (import.meta.env.VITE_WS_URL) {
|
|
return import.meta.env.VITE_WS_URL;
|
|
}
|
|
|
|
const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
|
|
return `${protocol}//${window.location.host}/ws`;
|
|
}
|
|
|
|
export function useWebSocket({ onMessage }) {
|
|
const onMessageRef = useRef(onMessage);
|
|
onMessageRef.current = onMessage;
|
|
|
|
useEffect(() => {
|
|
let socket;
|
|
let reconnectTimer;
|
|
let progressFlushTimer = null;
|
|
let progressFlushRaf = null;
|
|
let latestProgressMessage = null;
|
|
let isUnmounted = false;
|
|
|
|
const clearProgressFlush = () => {
|
|
if (progressFlushRaf !== null && typeof window !== 'undefined' && typeof window.cancelAnimationFrame === 'function') {
|
|
window.cancelAnimationFrame(progressFlushRaf);
|
|
}
|
|
if (progressFlushTimer) {
|
|
clearTimeout(progressFlushTimer);
|
|
}
|
|
progressFlushRaf = null;
|
|
progressFlushTimer = null;
|
|
};
|
|
|
|
const flushProgressMessage = () => {
|
|
progressFlushRaf = null;
|
|
progressFlushTimer = null;
|
|
const message = latestProgressMessage;
|
|
latestProgressMessage = null;
|
|
if (!message) {
|
|
return;
|
|
}
|
|
onMessageRef.current?.(message);
|
|
};
|
|
|
|
const scheduleProgressFlush = () => {
|
|
if (progressFlushRaf !== null || progressFlushTimer) {
|
|
return;
|
|
}
|
|
if (typeof window !== 'undefined' && typeof window.requestAnimationFrame === 'function') {
|
|
progressFlushRaf = window.requestAnimationFrame(flushProgressMessage);
|
|
return;
|
|
}
|
|
progressFlushTimer = setTimeout(flushProgressMessage, 16);
|
|
};
|
|
|
|
const connect = () => {
|
|
if (isUnmounted) {
|
|
return;
|
|
}
|
|
|
|
socket = new WebSocket(buildWsUrl());
|
|
|
|
socket.onmessage = (event) => {
|
|
try {
|
|
const message = JSON.parse(event.data);
|
|
if (message?.type === 'PIPELINE_PROGRESS') {
|
|
latestProgressMessage = message;
|
|
scheduleProgressFlush();
|
|
return;
|
|
}
|
|
onMessageRef.current?.(message);
|
|
} catch (error) {
|
|
// ignore invalid json
|
|
}
|
|
};
|
|
|
|
socket.onclose = () => {
|
|
if (!isUnmounted) {
|
|
reconnectTimer = setTimeout(connect, 1500);
|
|
}
|
|
};
|
|
|
|
socket.onerror = () => {
|
|
if (socket && socket.readyState !== WebSocket.CLOSED) {
|
|
socket.close();
|
|
}
|
|
};
|
|
};
|
|
|
|
connect();
|
|
|
|
return () => {
|
|
isUnmounted = true;
|
|
clearProgressFlush();
|
|
if (reconnectTimer) {
|
|
clearTimeout(reconnectTimer);
|
|
}
|
|
if (socket && socket.readyState !== WebSocket.CLOSED) {
|
|
socket.close();
|
|
}
|
|
};
|
|
}, []);
|
|
}
|