Add websocket ping for detect connection lost for client

Signed-off-by: YoungSoo Shin <shinys000114@gmail.com>

main/service/ws.c

@@ -13,6 +13,7 @@
 #include "pb.h"
 #include "pb_encode.h"
 #include "status.pb.h"
+#include "string.h" // Added for strlen and strncmp
 #include "webserver.h"
 
 #define UART_NUM UART_NUM_1
@@ -264,7 +265,6 @@ static esp_err_t ws_handler(httpd_req_t* req)
     httpd_ws_frame_t ws_pkt = {0};
     uint8_t buf[BUF_SIZE];
     ws_pkt.payload = buf;
-    ws_pkt.type = HTTPD_WS_TYPE_BINARY;
 
     esp_err_t ret = httpd_ws_recv_frame(req, &ws_pkt, BUF_SIZE);
     if (ret != ESP_OK)
@@ -273,7 +273,33 @@ static esp_err_t ws_handler(httpd_req_t* req)
         return ret;
     }
 
-    uart_write_bytes(UART_NUM, (const char*)ws_pkt.payload, ws_pkt.len);
+    if (ws_pkt.type == HTTPD_WS_TYPE_TEXT && ws_pkt.len == strlen("ping") &&
+        strncmp((const char*)ws_pkt.payload, "ping", ws_pkt.len) == 0)
+    {
+        ESP_LOGD(TAG, "Received application-level ping from client, sending pong.");
+        httpd_ws_frame_t pong_pkt = {
+            .payload = (uint8_t*)"pong", .len = strlen("pong"), .type = HTTPD_WS_TYPE_TEXT, .final = true};
+        return httpd_ws_send_frame(req, &pong_pkt);
+    }
+    else if (ws_pkt.type == HTTPD_WS_TYPE_CLOSE)
+    {
+        ESP_LOGI(TAG, "Client sent close frame, closing connection.");
+        return ESP_OK;
+    }
+    else if (ws_pkt.type == HTTPD_WS_TYPE_PING)
+    {
+        ESP_LOGD(TAG, "Received WebSocket PING control frame (handled by httpd).");
+        return ESP_OK;
+    }
+    else if (ws_pkt.type == HTTPD_WS_TYPE_PONG)
+    {
+        ESP_LOGD(TAG, "Received WebSocket PONG control frame.");
+        return ESP_OK;
+    }
+    else
+    {
+        uart_write_bytes(UART_NUM, (const char*)ws_pkt.payload, ws_pkt.len);
+    }
 
     return ESP_OK;
 }
@@ -327,4 +353,4 @@ void push_data_to_ws(const uint8_t* data, size_t len)
     }
 }
 
-esp_err_t change_baud_rate(int baud_rate) { return uart_set_baudrate(UART_NUM, baud_rate); }
+esp_err_t change_baud_rate(int baud_rate) { return uart_set_baudrate(UART_NUM, baud_rate); }

page/src/websocket.js

@@ -2,7 +2,7 @@
  * @file websocket.js
  * @description This module handles the WebSocket connection for real-time, two-way
  * communication with the server. It provides functions to initialize the connection
- * and send messages.
+ * and send messages, including a heartbeat mechanism to detect disconnections.
  */
 
 // The WebSocket instance, exported for potential direct access if needed.
@@ -11,14 +11,58 @@ export let websocket;
 // The WebSocket server address, derived from the current page's host (hostname + port).
 const baseGateway = `ws://${window.location.host}/ws`;
 
+// Heartbeat related variables
+let pingIntervalId = null;
+let pongTimeoutId = null;
+const HEARTBEAT_INTERVAL = 10000; // 10 seconds: How often to send a 'ping'
+const HEARTBEAT_TIMEOUT = 5000; // 5 seconds: How long to wait for a 'pong' after sending a 'ping'
+
 /**
- * Initializes the WebSocket connection and sets up event handlers.
- * @param {Object} callbacks - An object containing callback functions for WebSocket events.
- * @param {function} callbacks.onOpen - Called when the connection is successfully opened.
- * @param {function} callbacks.onClose - Called when the connection is closed.
- * @param {function} callbacks.onMessage - Called when a message is received from the server.
+ * Starts the heartbeat mechanism.
+ * Sends a 'ping' message to the server at regular intervals and sets a timeout
+ * to detect if a 'pong' response is not received.
  */
-export function initWebSocket({onOpen, onClose, onMessage}) {
+function startHeartbeat() {
+    stopHeartbeat(); // Ensure any previous heartbeat is stopped before starting a new one
+
+    pingIntervalId = setInterval(() => {
+        if (websocket && websocket.readyState === WebSocket.OPEN) {
+            websocket.send('ping');
+            console.log('WebSocket: Ping sent.');
+
+            // Set a timeout to check if a pong is received within HEARTBEAT_TIMEOUT
+            pongTimeoutId = setTimeout(() => {
+                console.warn('WebSocket: No pong received within timeout, closing connection.');
+                // If no pong is received, close the connection. This will trigger the onClose handler.
+                websocket.close();
+            }, HEARTBEAT_TIMEOUT);
+        }
+    }, HEARTBEAT_INTERVAL);
+}
+
+/**
+ * Stops the heartbeat mechanism by clearing the ping interval and pong timeout.
+ */
+function stopHeartbeat() {
+    if (pingIntervalId) {
+        clearInterval(pingIntervalId);
+        pingIntervalId = null;
+    }
+    if (pongTimeoutId) {
+        clearTimeout(pongTimeoutId);
+        pongTimeoutId = null;
+    }
+}
+
+/**
+ * Initializes the WebSocket connection and sets up event handlers, including a heartbeat mechanism.
+ * @param {Object} callbacks - An object containing callback functions for WebSocket events.
+ * @param {function} [callbacks.onOpen] - Called when the connection is successfully opened.
+ * @param {function} [callbacks.onClose] - Called when the connection is closed.
+ * @param {function} [callbacks.onMessage] - Called when a message is received from the server (excluding 'pong' messages).
+ * @param {function} [callbacks.onError] - Called when an error occurs with the WebSocket connection.
+ */
+export function initWebSocket({onOpen, onClose, onMessage, onError}) {
     const token = localStorage.getItem('authToken');
     let gateway = baseGateway;
 
@@ -31,10 +75,45 @@ export function initWebSocket({onOpen, onClose, onMessage}) {
     // Set binary type to arraybuffer to handle raw binary data from the UART.
     websocket.binaryType = "arraybuffer";
 
-    // Assign event handlers from the provided callbacks
-    if (onOpen) websocket.onopen = onOpen;
-    if (onClose) websocket.onclose = onClose;
-    if (onMessage) websocket.onmessage = onMessage;
+    // Assign event handlers, wrapping user-provided callbacks to include heartbeat logic
+    websocket.onopen = (event) => {
+        console.log('WebSocket connection opened.');
+        startHeartbeat(); // Start heartbeat on successful connection
+        if (onOpen) {
+            onOpen(event);
+        }
+    };
+
+    websocket.onclose = (event) => {
+        console.log('WebSocket connection closed:', event);
+        stopHeartbeat(); // Stop heartbeat when connection closes
+        if (onClose) {
+            onClose(event);
+        }
+    };
+
+    websocket.onmessage = (event) => {
+        if (event.data === 'pong') {
+            console.log('WebSocket: Pong received.');
+            // Clear the timeout as pong was received, resetting for the next ping
+            clearTimeout(pongTimeoutId);
+            pongTimeoutId = null;
+        } else {
+            // If it's not a pong message, pass it to the user's onMessage callback
+            if (onMessage) {
+                onMessage(event);
+            } else {
+                console.log('WebSocket message received:', event.data);
+            }
+        }
+    };
+
+    websocket.onerror = (error) => {
+        console.error('WebSocket error:', error);
+        if (onError) {
+            onError(error);
+        }
+    };
 }
 
 /**
@@ -44,5 +123,7 @@ export function initWebSocket({onOpen, onClose, onMessage}) {
 export function sendWebsocketMessage(data) {
     if (websocket && websocket.readyState === WebSocket.OPEN) {
         websocket.send(data);
+    } else {
+        console.warn('WebSocket is not open. Message not sent:', data);
     }
-}
+}