From 721084923ae9886ba52fbee94ec0c123482a3354 Mon Sep 17 00:00:00 2001
From: dijunkun <junkun.di@hotmail.com>
Date: Thu, 21 Nov 2024 15:52:09 +0800
Subject: [PATCH] [fix] release hdl from transmisson when connection closed

---
 src/client_id_generator.h    |  1 +
 src/signal_server.cpp        | 54 ++++++++++++++++++++++++++++--------
 src/signal_server.h          |  2 ++
 src/transmission_manager.cpp |  4 +--
 4 files changed, 48 insertions(+), 13 deletions(-)

diff --git a/src/client_id_generator.h b/src/client_id_generator.h
index 044dbbc..a7bd4e2 100644
--- a/src/client_id_generator.h
+++ b/src/client_id_generator.h
@@ -8,6 +8,7 @@
 #define _CLIENT_ID_GENERATOR_H_
 
 #include <iostream>
+#include <string>
 
 class ClientIdGenerator {
  public:
diff --git a/src/signal_server.cpp b/src/signal_server.cpp
index e87bd90..5042477 100644
--- a/src/signal_server.cpp
+++ b/src/signal_server.cpp
@@ -29,6 +29,9 @@ SignalServer::SignalServer() {
   server_.set_close_handler(
       std::bind(&SignalServer::on_close, this, std::placeholders::_1));
 
+  server_.set_fail_handler(
+      std::bind(&SignalServer::on_fail, this, std::placeholders::_1));
+
   server_.set_message_handler(std::bind(&SignalServer::on_message, this,
                                         std::placeholders::_1,
                                         std::placeholders::_2));
@@ -53,22 +56,42 @@ bool SignalServer::on_close(websocketpp::connection_hdl hdl) {
     LOG_INFO("Websocket connection [{}|{}] closed", ws_connections_[hdl],
              user_id);
 
-    std::string transmission_id = transmission_manager_.IsHost(user_id);
+    // check user is host or not
+    std::string transmission_id_host = transmission_manager_.IsHost(user_id);
+    if (!transmission_id_host.empty()) {
+      transmission_manager_.ReleaseTransmission(transmission_id_host);
+      LOG_INFO("Release transmission [{}] due to host [{}] leaves",
+               transmission_id_host, user_id);
 
-    if (!transmission_id.empty()) {
-      transmission_manager_.ReleaseTransmission(transmission_id);
-      LOG_INFO("Release transmission [{}] due to host leaves", transmission_id);
-    } else {
-      transmission_id = transmission_manager_.IsGuest(user_id);
-    }
-
-    if (!transmission_id.empty()) {
+      // notify all users in transmission
       json message = {{"type", "user_leave_transmission"},
-                      {"transmission_id", transmission_id},
+                      {"transmission_id", transmission_id_host},
                       {"user_id", user_id}};
 
       std::vector<std::string> user_id_list =
-          transmission_manager_.GetAllUserIdOfTransmission(transmission_id);
+          transmission_manager_.GetAllUserIdOfTransmission(
+              transmission_id_host);
+
+      for (const auto& user_id : user_id_list) {
+        send_msg(transmission_manager_.GetWsHandle(user_id), message);
+      }
+    }
+
+    // check user is guest or not
+    std::string transmission_id_guest = transmission_manager_.IsGuest(user_id);
+    if (!transmission_id_guest.empty()) {
+      transmission_manager_.ReleaseGuestFromTransmission(user_id);
+      LOG_INFO("Release guest [{}] from transmission [{}]", user_id,
+               transmission_id_guest);
+
+      // notify all users in transmission
+      json message = {{"type", "user_leave_transmission"},
+                      {"transmission_id", transmission_id_guest},
+                      {"user_id", user_id}};
+
+      std::vector<std::string> user_id_list =
+          transmission_manager_.GetAllUserIdOfTransmission(
+              transmission_id_guest);
 
       for (const auto& user_id : user_id_list) {
         send_msg(transmission_manager_.GetWsHandle(user_id), message);
@@ -81,6 +104,15 @@ bool SignalServer::on_close(websocketpp::connection_hdl hdl) {
   return true;
 }
 
+bool SignalServer::on_fail(websocketpp::connection_hdl hdl) {
+  std::string user_id = transmission_manager_.GetUserId(hdl);
+  if (!user_id.empty()) {
+    LOG_INFO("Websocket connection [{}|{}] failed", ws_connections_[hdl],
+             user_id);
+  }
+  return true;
+}
+
 bool SignalServer::on_ping(websocketpp::connection_hdl hdl, std::string s) {
   transmission_manager_.UpdateWsHandleLastActiveTime(hdl);
   return true;
diff --git a/src/signal_server.h b/src/signal_server.h
index 0ed6282..1a5fa20 100644
--- a/src/signal_server.h
+++ b/src/signal_server.h
@@ -27,6 +27,8 @@ class SignalServer {
 
   bool on_close(websocketpp::connection_hdl hdl);
 
+  bool on_fail(websocketpp::connection_hdl hdl);
+
   bool on_ping(websocketpp::connection_hdl hdl, std::string s);
 
   bool on_pong(websocketpp::connection_hdl hdl, std::string s);
diff --git a/src/transmission_manager.cpp b/src/transmission_manager.cpp
index 3aaf886..4e1c619 100644
--- a/src/transmission_manager.cpp
+++ b/src/transmission_manager.cpp
@@ -206,8 +206,8 @@ bool TransmissionManager::ReleaseGuestFromTransmission(
     auto guest_id_it =
         std::find(guest_id_list.begin(), guest_id_list.end(), guest_id);
     if (guest_id_it != guest_id_list.end()) {
-      LOG_INFO("Remove guest id [{}] from transmission [{}]", guest_id,
-               trans_it->first);
+      // LOG_INFO("Remove guest id [{}] from transmission [{}]", guest_id,
+      //          trans_it->first);
       guest_id_list.erase(guest_id_it);
       return true;
     }