import { ref } from 'vue';
/**
- Custom WebSocket hook that provides functionality to connect, send messages,
- and handle reconnection with heartbeat support.
- @param wsUri - The WebSocket URI to connect to.
- @param heartbeatMessage - The message to be sent periodically as a heartbeat. Default is 'heartbeat'.
- @param reconnectInterval - Interval (in ms) to attempt reconnection after WebSocket disconnection. Default is 5000 ms.
- @param timeout - Interval (in ms) between heartbeats and connection checks. Default is 20000 ms.
- @returns - An object containing the WebSocket reference,
methods to send messages, initialize connection, and close connection manually.
*/
export function useWebSocket(wsUri: string, heartbeatMessage = 'heartbeat', reconnectInterval = 5000, timeout = 20000) {
// Reactive reference for the WebSocket connection
const websock = ref<WebSocket | null>(null);
// Boolean flag to avoid multiple reconnections simultaneously
const lockReconnect = ref<boolean>(false);
// Variables to hold the timeouts for heartbeat and reconnection logic
let timeoutObj: ReturnType<typeof setTimeout> | null = null;
let serverTimeoutObj: ReturnType<typeof setTimeout> | null = null;
let timeoutnum: ReturnType<typeof setTimeout> | null = null;
/**
- Initialize the WebSocket connection.
- Sets up the event handlers for open, error, message, and close events.
*/
const initWebSocket = () => {
if (websock.value) {
console.log('WebSocket is already initialized.');
return;
}
// Create a new WebSocket connection
websock.value = new WebSocket(wsUri);
// Set up event handlers for WebSocket
websock.value.onopen = websocketonopen;
websock.value.onerror = websocketonerror;
websock.value.onmessage = websocketonmessage;
websock.value.onclose = websocketclose;
};
/**
- Event handler for WebSocket 'onopen' event.
- This is triggered when the WebSocket connection is successfully opened.
*/
const websocketonopen = () => {
websocketsend('Connection established'); // Send a confirmation message
startHeartbeat(); // Start sending heartbeat messages
};
/**
- Event handler for WebSocket 'onerror' event.
- This is triggered when there's an error in the WebSocket connection.
- @param e - The event object from the error.
*/
const websocketonerror = (e: Event) => {
console.error('WebSocket connection error', e);
websock.value = null;
reconnect(); // Try to reconnect when an error occurs
};
/**
- Event handler for WebSocket 'onmessage' event.
- This is triggered when the WebSocket receives a message.
- @param event - The message event object from the WebSocket.
*/
const websocketonmessage = (event: MessageEvent) => {
console.log('Received message from server:', event.data);
resetHeartbeat(); // Reset the heartbeat timer after receiving a message
// Handle the received data
try {
const data = JSON.parse(event.data); // Parse the JSON message from the server
// Process the data as needed (you can customize this part)
console.log('Parsed data:', data);
} catch (e) {
console.error('Error parsing JSON:', e); // Log any errors with JSON parsing
}
};
/**
- Event handler for WebSocket 'onclose' event.
- This is triggered when the WebSocket connection is closed.
- @param e - The close event object.
*/
const websocketclose = (e: CloseEvent) => {
console.log('WebSocket connection closed:', e);
reconnect(); // Attempt to reconnect when the WebSocket is closed
};
/**
- Send a message through the WebSocket connection.
- @param msg - The message to be sent through the WebSocket.
*/
const websocketsend = (msg: string) => {
if (websock.value?.readyState === WebSocket.OPEN) {
websock.value.send(msg); // Send the message if the WebSocket is open
}
};
/**
- Reconnect the WebSocket connection.
- This is triggered when the WebSocket connection fails and needs to reconnect.
*/
const reconnect = () => {
if (lockReconnect.value) return; // If reconnecting is locked, do nothing
lockReconnect.value = true; // Lock the reconnection to prevent multiple attempts
// Clear any existing reconnection timeouts
timeoutnum && clearTimeout(timeoutnum);
// Attempt reconnection after a delay
timeoutnum = setTimeout(() => {
initWebSocket(); // Re-initialize WebSocket connection
lockReconnect.value = false; // Unlock reconnection
}, reconnectInterval); // Reconnect after specified interval
};
/**
- Reset the heartbeat mechanism.
- Called when a message is received to reset the heartbeat timers.
*/
const resetHeartbeat = () => {
timeoutObj && clearTimeout(timeoutObj); // Clear the current heartbeat timeout
serverTimeoutObj && clearTimeout(serverTimeoutObj); // Clear the server-side timeout
startHeartbeat(); // Start the heartbeat mechanism again
};
/**
-
Start the heartbeat mechanism.
-
Periodically sends a heartbeat message to the server to keep the connection alive.
*/
const startHeartbeat = () => {
timeoutObj && clearTimeout(timeoutObj); // Clear existing heartbeat timeouts
serverTimeoutObj && clearTimeout(serverTimeoutObj);
console.log('心跳');
// Start the timeout to send a heartbeat after the specified interval
timeoutObj = setTimeout(() => {
// If the WebSocket connection is open, send a heartbeat message
if (websock.value?.readyState === WebSocket.OPEN) {
websocketsend(heartbeatMessage); // Send the heartbeat message
} else {
reconnect(); // Reconnect if the WebSocket is not open
}// Set another timeout to close the connection if no heartbeat response is received
serverTimeoutObj = setTimeout(() => {
websock.value?.close(); // Close the WebSocket if no response is received
}, timeout);
}, timeout);
};
/**
- Manually close the WebSocket connection.
- This clears any ongoing timeouts and closes the WebSocket.
*/
const closeWebSocket = () => {
if (websock.value) {
websock.value.close(); // Close the WebSocket connection
websock.value = null; // Set WebSocket reference to null
}
// Clear all timeouts related to heartbeat and reconnection
timeoutObj && clearTimeout(timeoutObj);
serverTimeoutObj && clearTimeout(serverTimeoutObj);
timeoutnum && clearTimeout(timeoutnum);
console.log('WebSocket connection closed manually');
};
// Automatically initialize WebSocket on component mount
// onMounted(() => {
// initWebSocket();
// });
// // Automatically close WebSocket on component unmount to prevent memory leaks
// onBeforeUnmount(() => {
// closeWebSocket();
// });
return {
websock, // Expose WebSocket reference
websocketsend, // Expose method to send messages
initWebSocket, // Expose method to manually initialize WebSocket
closeWebSocket // Expose method to manually close WebSocket
};
}
使用
import { useWebSocket } from './useWebSocket'; // Import the WebSocket hook (adjust path as necessary)
const wsUri = ‘wss://’
const { websock, websocketsend, initWebSocket, closeWebSocket } = useWebSocket(wsUri);
网友评论