Fix memory leak when client crashed (#1968)

Fixes #1967 Signed-off-by: vegetableysm <[email protected]>
v6d-io · Jul 31, 2024 · a0ad12c · a0ad12c
1 parent 388fcaf
commit a0ad12c
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 4 deletions.
diff --git a/src/common/rdma/rdma.h b/src/common/rdma/rdma.h
@@ -77,6 +77,7 @@ class IRDMA {
     for (auto& res : vec) {
       RETURN_ON_ERROR(CloseResource(res, resource_name));
     }
+    vec.clear();
     return Status::OK();
   }
 
@@ -86,6 +87,7 @@ class IRDMA {
     for (auto iter = mapping.begin(); iter != mapping.end(); ++iter) {
       RETURN_ON_ERROR(CloseResource(iter->second, resource_name));
     }
+    mapping.clear();
     return Status::OK();
   }
 

diff --git a/src/common/rdma/rdma_server.cc b/src/common/rdma/rdma_server.cc
@@ -145,6 +145,12 @@ Status RDMAServer::Close() {
     RETURN_ON_ERROR(CloseResourcesInMap(ep_map_, "endpoint created by server"));
   }
 
+  {
+    std::lock_guard<std::mutex> lock(wait_conn_ep_map_mutex_);
+    RETURN_ON_ERROR(
+        CloseResourcesInMap(wait_conn_ep_map_, "endpoint created by server"));
+  }
+
   RETURN_ON_ERROR(CloseResource(txcq, "transmit comeple queue"));
   RETURN_ON_ERROR(CloseResource(rxcq, "receive comeple queue"));
   RETURN_ON_ERROR(CloseResource(pep, "passive endpoint"));
@@ -240,7 +246,10 @@ Status RDMAServer::AddClient(uint64_t& ep_token, void* ep) {
 
 Status RDMAServer::RemoveClient(uint64_t ep_token) {
   std::lock_guard<std::mutex> lock(ep_map_mutex_);
-  ep_map_.erase(ep_token);
+  if (ep_map_.find(ep_token) != ep_map_.end()) {
+    CloseResource(ep_map_[ep_token], "client endpoint");
+    ep_map_.erase(ep_token);
+  }
   return Status::OK();
 }
 
@@ -249,6 +258,7 @@ Status RDMAServer::RemoveClient(fid_ep* ep) {
   for (auto iter = ep_map_.begin(); iter != ep_map_.end(); iter++) {
     if (iter->second == ep) {
       ep_map_.erase(iter);
+      CloseResource(ep, "client endpoint");
       return Status::OK();
     }
   }
@@ -272,11 +282,13 @@ Status RDMAServer::RegisterMemory(fid_mr** mr, void* address, size_t size,
 }
 
 Status RDMAServer::DeregisterMemory(RegisterMemInfo& memInfo) {
+  {
+    std::lock_guard<std::mutex> lock(mr_array_mutex_);
+    mr_array.erase(std::remove(mr_array.begin(), mr_array.end(), memInfo.mr),
+                   mr_array.end());
+  }
   VINEYARD_CHECK_OK(IRDMA::CloseResource(reinterpret_cast<fid_mr*>(memInfo.mr),
                                          "memory region"));
-  std::lock_guard<std::mutex> lock(mr_array_mutex_);
-  mr_array.erase(std::remove(mr_array.begin(), mr_array.end(), memInfo.mr),
-                 mr_array.end());
   return Status::OK();
 }