68 lines
1.6 KiB
JavaScript
68 lines
1.6 KiB
JavaScript
import { useState, useEffect, useRef, useCallback } from 'react';
|
|
|
|
export const WebSocketState = {
|
|
CONNECTING: 'CONNECTING',
|
|
CONNECTED: 'CONNECTED',
|
|
DISCONNECTED: 'DISCONNECTED',
|
|
};
|
|
|
|
export const useWebSocket = (url) => {
|
|
const [wsState, setWsState] = useState(WebSocketState.DISCONNECTED);
|
|
const wsRef = useRef(null);
|
|
const handlersRef = useRef([]);
|
|
|
|
const connect = useCallback(() => {
|
|
setWsState(WebSocketState.CONNECTING);
|
|
const ws = new WebSocket(url);
|
|
wsRef.current = ws;
|
|
|
|
ws.onopen = () => {
|
|
setWsState(WebSocketState.CONNECTED);
|
|
};
|
|
|
|
ws.onclose = () => {
|
|
setWsState(WebSocketState.DISCONNECTED);
|
|
setTimeout(connect, 2000);
|
|
};
|
|
|
|
ws.onerror = (error) => {
|
|
console.error('WebSocket error:', error);
|
|
};
|
|
|
|
ws.onmessage = (event) => {
|
|
try {
|
|
const message = JSON.parse(event.data);
|
|
handlersRef.current.forEach(handler => handler(message));
|
|
} catch (error) {
|
|
console.error('Error handling WebSocket message:', error);
|
|
}
|
|
};
|
|
|
|
return ws;
|
|
}, [url]);
|
|
|
|
useEffect(() => {
|
|
const ws = connect();
|
|
return () => {
|
|
if (ws) ws.close();
|
|
};
|
|
}, [connect]);
|
|
|
|
const sendMessage = useCallback((data) => {
|
|
if (wsRef.current?.readyState === WebSocket.OPEN) {
|
|
wsRef.current.send(JSON.stringify(data));
|
|
} else {
|
|
console.error('WebSocket not connected, state:', wsRef.current?.readyState);
|
|
}
|
|
}, []);
|
|
|
|
const addMessageHandler = useCallback((handler) => {
|
|
handlersRef.current.push(handler);
|
|
}, []);
|
|
|
|
return {
|
|
wsState,
|
|
sendMessage,
|
|
addMessageHandler
|
|
};
|
|
}; |