Implementation for websocket heartbeat

src/ws/ws_core.cpp

@@ -29,12 +29,16 @@ WsCore::~WsCore() {
   m_endpoint_.close(connection_handle_, websocketpp::close::status::going_away,
                     "", ec);
   if (ec) {
-    LOG_INFO("> Error closing connection {}", ec.message());
+    LOG_INFO("Closing connection error: {}", ec.message());
   }
 
   if (m_thread_->joinable()) {
     m_thread_->join();
   }
+
+  if (ping_thread_->joinable()) {
+    ping_thread_->join();
+  }
 }
 
 int WsCore::Connect(std::string const &uri) {
@@ -45,7 +49,7 @@ int WsCore::Connect(std::string const &uri) {
   connection_handle_ = con->get_handle();
 
   if (ec) {
-    LOG_INFO("> Connect initialization error: {}", ec.message());
+    LOG_INFO("Connect initialization error: {}", ec.message());
     return -1;
   }
 
@@ -57,18 +61,15 @@ int WsCore::Connect(std::string const &uri) {
       websocketpp::lib::bind(&WsCore::OnClose, this, &m_endpoint_,
                              websocketpp::lib::placeholders::_1));
 
-  // con->set_ping_handler(websocketpp::lib::bind(
-  //     &WsCore::on_ping,
-  //     this,
-  //     websocketpp::lib::placeholders::_1,
-  //     websocketpp::lib::placeholders::_2
-  // ));
+  con->set_ping_handler(websocketpp::lib::bind(
+      &WsCore::OnPing, this, websocketpp::lib::placeholders::_1,
+      websocketpp::lib::placeholders::_2));
 
   con->set_pong_handler(websocketpp::lib::bind(
       &WsCore::OnPong, this, websocketpp::lib::placeholders::_1,
       websocketpp::lib::placeholders::_2));
 
-  con->set_pong_timeout(1000);
+  con->set_pong_timeout(10);
 
   con->set_pong_timeout_handler(websocketpp::lib::bind(
       &WsCore::OnPongTimeout, this, websocketpp::lib::placeholders::_1,
@@ -88,7 +89,7 @@ void WsCore::Close(websocketpp::close::status::value code, std::string reason) {
 
   m_endpoint_.close(connection_handle_, code, reason, ec);
   if (ec) {
-    LOG_INFO("> Error initiating close: {}", ec.message());
+    LOG_INFO("Initiating close error: {}", ec.message());
   }
 }
 
@@ -98,20 +99,22 @@ void WsCore::Send(std::string message) {
   m_endpoint_.send(connection_handle_, message,
                    websocketpp::frame::opcode::text, ec);
   if (ec) {
-    LOG_INFO("> Error sending message: {}", ec.message());
+    LOG_INFO("Sending message error: {}", ec.message());
     return;
   }
 }
 
-void WsCore::Ping() {
-  websocketpp::lib::error_code ec;
+void WsCore::Ping(websocketpp::connection_hdl hdl) {
+  auto con = m_endpoint_.get_con_from_hdl(hdl);
+  while (con->get_state() == websocketpp::session::state::open) {
+    websocketpp::lib::error_code ec;
+    m_endpoint_.ping(hdl, "", ec);
+    if (ec) {
+      LOG_ERROR("Ping error: {}", ec.message());
+      break;
+    }
 
-  std::string message = "ping";
-
-  m_endpoint_.ping(connection_handle_, message, ec);
-  if (ec) {
-    LOG_INFO("> Error sending ping");
-    return;
+    std::this_thread::sleep_for(std::chrono::seconds(3));
   }
 }
 
@@ -119,6 +122,9 @@ const std::string &WsCore::GetStatus() { return connection_status_; }
 
 void WsCore::OnOpen(client *c, websocketpp::connection_hdl hdl) {
   connection_status_ = "Open";
+
+  ping_thread_ = websocketpp::lib::make_shared<websocketpp::lib::thread>(
+      &WsCore::Ping, this, hdl);
 }
 
 void WsCore::OnFail(client *c, websocketpp::connection_hdl hdl) {
@@ -129,9 +135,25 @@ void WsCore::OnClose(client *c, websocketpp::connection_hdl hdl) {
   connection_status_ = "Closed";
 }
 
-void WsCore::OnPong(websocketpp::connection_hdl, std::string msg) {}
+bool WsCore::OnPing(websocketpp::connection_hdl hdl, std::string msg) {
+  return true;
+}
 
-void WsCore::OnPongTimeout(websocketpp::connection_hdl, std::string msg) {}
+bool WsCore::OnPong(websocketpp::connection_hdl hdl, std::string msg) {
+  return true;
+}
+
+void WsCore::OnPongTimeout(websocketpp::connection_hdl hdl, std::string msg) {
+  if (timeout_count_ < 2) {
+    timeout_count_++;
+    return;
+  }
+
+  LOG_WARN("Pong timeout, reset connection");
+  // m_endpoint_.close(hdl, websocketpp::close::status::normal,
+  // "OnPongTimeout");
+  m_endpoint_.reset();
+}
 
 void WsCore::OnMessage(websocketpp::connection_hdl, client::message_ptr msg) {
   OnReceiveMessage(msg->get_payload());

src/ws/ws_core.h

@@ -4,6 +4,7 @@
 #include <map>
 #include <sstream>
 #include <string>
+#include <thread>
 
 #include "websocketpp/client.hpp"
 #include "websocketpp/common/memory.hpp"
@@ -24,7 +25,7 @@ class WsCore {
 
   void Send(std::string message);
 
-  void Ping();
+  void Ping(websocketpp::connection_hdl hdl);
 
   const std::string &GetStatus();
 
@@ -35,11 +36,13 @@ class WsCore {
 
   void OnClose(client *c, websocketpp::connection_hdl hdl);
 
-  void OnPong(websocketpp::connection_hdl, std::string msg);
+  bool OnPing(websocketpp::connection_hdl hdl, std::string msg);
 
-  void OnPongTimeout(websocketpp::connection_hdl, std::string msg);
+  bool OnPong(websocketpp::connection_hdl hdl, std::string msg);
 
-  void OnMessage(websocketpp::connection_hdl, client::message_ptr msg);
+  void OnPongTimeout(websocketpp::connection_hdl hdl, std::string msg);
+
+  void OnMessage(websocketpp::connection_hdl hdl, client::message_ptr msg);
 
   virtual void OnReceiveMessage(const std::string &msg) = 0;
 
@@ -47,8 +50,10 @@ class WsCore {
   client m_endpoint_;
   websocketpp::connection_hdl connection_handle_;
   websocketpp::lib::shared_ptr<websocketpp::lib::thread> m_thread_;
+  websocketpp::lib::shared_ptr<websocketpp::lib::thread> ping_thread_;
 
   std::string connection_status_ = "Connecting";
+  int timeout_count_ = 0;
 };
 
 #endif