1 files changed, 36 insertions, 23 deletions
diff --git a/src/components/transport_manager/test/sample_websocket_server.cc b/src/components/transport_manager/test/sample_websocket_server.cc
index 917184a361..9b0e5b59d5 100644
--- a/src/components/transport_manager/test/sample_websocket_server.cc
+++ b/src/components/transport_manager/test/sample_websocket_server.cc
@@ -51,11 +51,12 @@ void WSSession::WSServer::AddURLRoute(const std::string& target) {
 
 void WSSession::WSServer::Run() {
   req_ = {};
-  http::async_read(
-      ws_.next_layer(),
-      buffer_,
-      req_,
-      std::bind(&WSServer::OnWebsocketHandshake, this, std::placeholders::_1));
+  http::async_read(ws_.next_layer(),
+                   buffer_,
+                   req_,
+                   std::bind(&WSServer::OnWebsocketHandshake,
+                             shared_from_this(),
+                             std::placeholders::_1));
 }
 
 void WSSession::WSServer::OnWebsocketHandshake(
@@ -75,9 +76,10 @@ void WSSession::WSServer::OnWebsocketHandshake(
     // Accept the websocket handshake
     ws_.async_accept(
         req_,
-        boost::asio::bind_executor(
-            strand_,
-            std::bind(&WSServer::OnAccept, this, std::placeholders::_1)));
+        boost::asio::bind_executor(strand_,
+                                   std::bind(&WSServer::OnAccept,
+                                             shared_from_this(),
+                                             std::placeholders::_1)));
   }
 }
 
@@ -112,7 +114,8 @@ std::string WSSession::WSServer::ParseRouteFromTarget(
 }
 
 WSSession::WSSession(const std::string& address, uint16_t port)
-    : address_(address)
+    : io_pool_(1)
+    , address_(address)
     , port_(port)
     , acceptor_(ioc_)
     , socket_(ioc_)
@@ -152,8 +155,10 @@ WSSession::WSSession(const std::string& address, uint16_t port)
 void WSSession::Run() {
   if (acceptor_.is_open()) {
     acceptor_.async_accept(
-        socket_, std::bind(&WSSession::on_accept, this, std::placeholders::_1));
-    ioc_.run();
+        socket_,
+        std::bind(
+            &WSSession::on_accept, shared_from_this(), std::placeholders::_1));
+    boost::asio::post(io_pool_, [&]() { ioc_.run(); });
   }
 }
 
@@ -161,6 +166,7 @@ void WSSession::Stop() {
   try {
     ioc_.stop();
     acceptor_.close();
+    io_pool_.join();
   } catch (...) {
     std::cerr << "Failed to close connection" << std::endl;
   }
@@ -199,9 +205,10 @@ void WSSSession::WSSServer::AddURLRoute(const std::string& target) {
 }
 void WSSSession::WSSServer::Run() {
   // Perform the SSL handshake
-  wss_.next_layer().async_handshake(
-      ssl::stream_base::server,
-      std::bind(&WSSServer::OnSSLHandshake, this, std::placeholders::_1));
+  wss_.next_layer().async_handshake(ssl::stream_base::server,
+                                    std::bind(&WSSServer::OnSSLHandshake,
+                                              shared_from_this(),
+                                              std::placeholders::_1));
 }
 
 void WSSSession::WSSServer::OnSSLHandshake(beast::error_code ec) {
@@ -210,11 +217,12 @@ void WSSSession::WSSServer::OnSSLHandshake(beast::error_code ec) {
   }
 
   req_ = {};
-  http::async_read(
-      wss_.next_layer(),
-      buffer_,
-      req_,
-      std::bind(&WSSServer::OnWebsocketHandshake, this, std::placeholders::_1));
+  http::async_read(wss_.next_layer(),
+                   buffer_,
+                   req_,
+                   std::bind(&WSSServer::OnWebsocketHandshake,
+                             shared_from_this(),
+                             std::placeholders::_1));
 }
 
 void WSSSession::WSSServer::OnWebsocketHandshake(
@@ -233,7 +241,9 @@ void WSSSession::WSSServer::OnWebsocketHandshake(
     }
     // Accept the websocket handshake
     wss_.async_accept(
-        req_, std::bind(&WSSServer::OnAccept, this, std::placeholders::_1));
+        req_,
+        std::bind(
+            &WSSServer::OnAccept, shared_from_this(), std::placeholders::_1));
   }
 }
 
@@ -271,7 +281,8 @@ WSSSession::WSSSession(const std::string& address,
                        uint16_t port,
                        const std::string& certificate,
                        const std::string& private_key)
-    : acceptor_(ioc_)
+    : io_pool_(1)
+    , acceptor_(ioc_)
     , socket_(ioc_)
     , ctx_(ssl::context::sslv23_server)
     , wss_(nullptr) {
@@ -336,6 +347,7 @@ void WSSSession::Stop() {
   try {
     ioc_.stop();
     acceptor_.close();
+    io_pool_.join();
   } catch (...) {
     std::cerr << "Failed to close connection" << std::endl;
   }
@@ -353,8 +365,9 @@ void WSSSession::do_accept() {
   if (acceptor_.is_open()) {
     acceptor_.async_accept(
         socket_,
-        std::bind(&WSSSession::on_accept, this, std::placeholders::_1));
-    ioc_.run();
+        std::bind(
+            &WSSSession::on_accept, shared_from_this(), std::placeholders::_1));
+    boost::asio::post(io_pool_, [&]() { ioc_.run(); });
   }
 }