Fix destructor function bug of RPCServer (#1951)

Signed-off-by: vegetableysm <yuanshumin.ysm@alibaba-inc.com>
v6d-io · Jul 22, 2024 · ebe8f07 · ebe8f07
1 parent a561b47
commit ebe8f07
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 5 deletions.
diff --git a/src/server/async/rpc_server.cc b/src/server/async/rpc_server.cc
@@ -42,10 +42,6 @@ RPCServer::RPCServer(std::shared_ptr<VineyardServer> vs_ptr)
 }
 
 RPCServer::~RPCServer() {
-  if (acceptor_.is_open()) {
-    acceptor_.close();
-  }
-
   if (rdma_stop_) {
     return;
   }
@@ -64,6 +60,23 @@ RPCServer::~RPCServer() {
   rdma_stop_ = true;
 }
 
+void RPCServer::Stop() {
+  SocketServer::Stop();
+  boost::system::error_code ec;
+  ec = acceptor_.cancel(ec);
+  if (ec) {
+    LOG(ERROR) << "Failed to close the RPC server: " << ec.message();
+  }
+  while (true) {
+    {
+      std::lock_guard<std::mutex> scope_lock(accept_mutex_);
+      if (!is_accepting_) {
+        break;
+      }
+    }
+  }
+}
+
 Status RPCServer::InitRDMA() {
   std::string rdma_endpoint = RDMAEndpoint();
   size_t pos = rdma_endpoint.find(':');
@@ -123,6 +136,10 @@ void RPCServer::doAccept() {
   if (!acceptor_.is_open()) {
     return;
   }
+  {
+    std::lock_guard<std::mutex> scope_lock(accept_mutex_);
+    is_accepting_ = true;
+  }
   auto self(shared_from_this());
   acceptor_.async_accept(socket_, [self](boost::system::error_code ec) {
     if (!ec) {
@@ -143,7 +160,13 @@ void RPCServer::doAccept() {
     if (!ec || ec != boost::system::errc::operation_canceled) {
       if (!self->stopped_.load() || !self->closable_.load()) {
         self->doAccept();
+      } else {
+        std::lock_guard<std::mutex> scope_lock(self->accept_mutex_);
+        self->is_accepting_ = false;
       }
+    } else {
+      std::lock_guard<std::mutex> scope_lock(self->accept_mutex_);
+      self->is_accepting_ = false;
     }
   });
 }

diff --git a/src/server/async/rpc_server.h b/src/server/async/rpc_server.h
@@ -45,6 +45,8 @@ class RPCServer : public SocketServer,
 
   void Start() override;
 
+  void Stop() override;
+
   std::string Endpoint() {
     return get_hostname() + ":" + json_to_string(rpc_spec_["port"]);
   }
@@ -86,6 +88,8 @@ class RPCServer : public SocketServer,
   const json rpc_spec_;
   asio::ip::tcp::acceptor acceptor_;
   asio::ip::tcp::socket socket_;
+  std::mutex accept_mutex_;
+  bool is_accepting_ = false;
 
   // connection id to rdma server
   std::unordered_map<uint64_t, std::map<void*, RegisterMemInfo>>

diff --git a/src/server/async/socket_server.h b/src/server/async/socket_server.h
@@ -236,7 +236,7 @@ class SocketServer {
   /**
    * Call "Stop" on all connections, then clear the connection pool.
    */
-  void Stop();
+  virtual void Stop();
 
   /**
    * Cancel the "async_accept" action on the acceptor to stop accepting