implementing heartbeat

Author: sovcik

src/WebSockets.cpp

@@ -433,10 +433,12 @@ void WebSockets::handleWebsocketPayloadCb(WSclient_t * client, bool ok, uint8_t
                 break;
             case WSop_ping:
                 // send pong back
+                DEBUG_WEBSOCKETS("[WS][%d][handleWebsocket] ping received (%s)\n", client->num, payload ? (const char*)payload : "");
                 sendFrame(client, WSop_pong, payload, header->payloadLen);
                 break;
             case WSop_pong:
                 DEBUG_WEBSOCKETS("[WS][%d][handleWebsocket] get pong (%s)\n", client->num, payload ? (const char*)payload : "");
+                client->pongReceived = true;
                 break;
             case WSop_close: {
                 #ifndef NODEBUG_WEBSOCKETS
@@ -652,3 +654,47 @@ size_t WebSockets::write(WSclient_t * client, const char *out) {
 	if(out == NULL) return 0;
 	return write(client, (uint8_t*)out, strlen(out));
 }
+
+/**
+ * enable ping/pong heartbeat process
+ * @param client WSclient_t *
+ * @param pingInterval uint32_t how often ping will be sent
+ * @param pongTimeout uint32_t millis after which pong should timout if not received
+ * @param disconnectTimeoutCount uint8_t how many timeouts before disconnect, 0=> do not disconnect
+ */
+void WebSockets::enableHeartbeat(WSclient_t * client, uint32_t pingInterval, uint32_t pongTimeout, uint8_t disconnectTimeoutCount){
+    if(client == NULL) return;
+    client->pingInterval = pingInterval;
+    client->pongTimeout = pongTimeout;
+    client->disconnectTimeoutCount = disconnectTimeoutCount;
+
+}
+
+/**
+ * handle ping/pong heartbeat timeout process
+ * @param client WSclient_t *
+ */ 
+void WebSockets::handleHBTimeout(WSclient_t * client){
+    if (client->pingInterval) { // if heartbeat is enabled
+        uint32_t pi = millis() - client->lastPing;
+
+        if (client->pongReceived) {
+            client->pongTimeoutCount = 0;
+        } else {
+            if (pi > client->pongTimeout){ // pong not received in time
+                client->pongTimeoutCount++;
+                client->lastPing = millis() - client->pingInterval - 500; // force ping on the next run 
+
+                DEBUG_WEBSOCKETS("[HBtimeout] pong TIMEOUT! lp=%d millis=%d pi=%d count=%d\n", client->lastPing, millis(), pi, client->pongTimeoutCount);
+
+                if (client->disconnectTimeoutCount && client->pongTimeoutCount >= client->disconnectTimeoutCount){
+                    DEBUG_WEBSOCKETS("[HBtimeout] count=%d, DISCONNECTING\n", client->pongTimeoutCount);
+                    client->status = WSC_NOT_CONNECTED;
+                    //clientDisconnect(client);
+                }
+            } 
+        }
+
+    }
+        
+}

src/WebSockets.h

@@ -264,6 +264,13 @@ typedef struct {
         bool cHttpHeadersValid; ///< non-websocket http header validity indicator
         size_t cMandatoryHeadersCount; ///< non-websocket mandatory http headers present count
 
+        bool pongReceived;
+        uint32_t pingInterval;  // how often ping will be sent, 0 means "heartbeat is not active"
+        uint32_t lastPing;      // millis when last pong has been received
+        uint32_t pongTimeout;   // interval in millis after which pong is considered to timeout
+        uint8_t disconnectTimeoutCount;  // after how many subsequent pong timeouts discconnect will happen, 0 means "do not disconnect"
+        uint8_t pongTimeoutCount;   // current pong timeout count     
+
 #if (WEBSOCKETS_NETWORK_TYPE == NETWORK_ESP8266_ASYNC)
         String cHttpLine;   ///< HTTP header lines
 #endif
@@ -303,6 +310,9 @@ class WebSockets {
         virtual size_t write(WSclient_t * client, uint8_t *out, size_t n);
         size_t write(WSclient_t * client, const char *out);
 
+        void enableHeartbeat(WSclient_t * client, uint32_t pingInterval, uint32_t pongTimeout, uint8_t disconnectTimeoutCount);
+        void handleHBTimeout(WSclient_t * client);
+
 
 };
 

src/WebSocketsClient.cpp

@@ -66,6 +66,10 @@ void WebSocketsClient::begin(const char *host, uint16_t port, const char * url,
     _client.plainAuthorization = "";
     _client.isSocketIO = false;
 
+    _client.lastPing = 0;
+    _client.pongReceived = false;
+    _client.pongTimeoutCount = 0;    
+
 #ifdef ESP8266
     randomSeed(RANDOM_REG32);
 #else
@@ -169,6 +173,10 @@ void WebSocketsClient::loop(void) {
 
         }
     } else {
+        if (_client.status == WSC_CONNECTED){
+            handleHBPing();
+            handleHBTimeout(&_client);
+        }
         handleClientData();
     }
 }
@@ -243,7 +251,10 @@ bool WebSocketsClient::sendBIN(const uint8_t * payload, size_t length) {
  */
 bool WebSocketsClient::sendPing(uint8_t * payload, size_t length) {
     if(clientIsConnected(&_client)) {
-        return sendFrame(&_client, WSop_ping, payload, length);
+        bool sent = sendFrame(&_client, WSop_ping, payload, length);
+        if (sent) 
+            _client.lastPing = millis();
+        return sent;
     }
     return false;
 }
@@ -761,3 +772,36 @@ void WebSocketsClient::asyncConnect() {
 }
 
 #endif
+
+/**
+ * send heartbeat ping to server in set intervals
+ */ 
+void WebSocketsClient::handleHBPing(){
+    if (_client.pingInterval == 0) return;
+    uint32_t pi = millis() - _client.lastPing;
+    if (pi > _client.pingInterval){
+        DEBUG_WEBSOCKETS("[WS-Client] sending HB ping\n");
+        if (sendPing()) {
+            _client.lastPing = millis();
+            _client.pongReceived = false;
+        }
+    }
+
+}
+
+/**
+ * enable ping/pong heartbeat process
+ * @param pingInterval uint32_t how often ping will be sent
+ * @param pongTimeout uint32_t millis after which pong should timout if not received
+ * @param disconnectTimeoutCount uint8_t how many timeouts before disconnect, 0=> do not disconnect
+ */
+void WebSocketsClient::enableHeartbeat(uint32_t pingInterval, uint32_t pongTimeout, uint8_t disconnectTimeoutCount){
+    WebSockets::enableHeartbeat(&_client, pingInterval, pongTimeout, disconnectTimeoutCount);
+}
+
+/**
+ * disable ping/pong heartbeat process
+ */
+void WebSocketsClient::disableHeartbeat(){
+    _client.pingInterval = 0;
+}

src/WebSocketsClient.h

@@ -86,6 +86,9 @@ class WebSocketsClient: private WebSockets {
 
         void setReconnectInterval(unsigned long time);
 
+        void enableHeartbeat(uint32_t pingInterval, uint32_t pongTimeout, uint8_t disconnectTimeoutCount);
+        void disableHeartbeat();        
+
     protected:
         String _host;
         uint16_t _port;
@@ -115,6 +118,8 @@ class WebSocketsClient: private WebSockets {
         void connectedCb();
         void connectFailedCb();
 
+        void handleHBPing(); // send ping in specified intervals
+
 #if (WEBSOCKETS_NETWORK_TYPE == NETWORK_ESP8266_ASYNC)
         void asyncConnect();
 #endif