From e7fe3bff2d00c5f30b133e5d5d95764d4be1f123 Mon Sep 17 00:00:00 2001
From: lijunlin <70465472+L-Xiafeng@users.noreply.github.com>
Date: Wed, 20 Nov 2024 12:58:42 +0800
Subject: [PATCH] fix: Fix sigsegv in cranectld (#368)

* fix: fix sigsegv in cranectld
m_pending_task_map_ entry not erased.

Signed-off-by: Li Junlin <xiafeng.li@foxmail.com>

* fix: Data race in ConnectCranedNode_
add lock to avoid data race

Signed-off-by: Li Junlin <xiafeng.li@foxmail.com>

* Refactor.

Signed-off-by: RileyW <wrllrwwrllrw@gmail.com>

---------

Signed-off-by: Li Junlin <xiafeng.li@foxmail.com>
Signed-off-by: RileyW <wrllrwwrllrw@gmail.com>
Co-authored-by: RileyW <wrllrwwrllrw@gmail.com>
---
 src/CraneCtld/CranedKeeper.cpp | 47 +++++++++++++++++++---------------
 src/CraneCtld/TaskScheduler.h  |  1 +
 2 files changed, 27 insertions(+), 21 deletions(-)

diff --git a/src/CraneCtld/CranedKeeper.cpp b/src/CraneCtld/CranedKeeper.cpp
index 55532b79b..4a90b715a 100644
--- a/src/CraneCtld/CranedKeeper.cpp
+++ b/src/CraneCtld/CranedKeeper.cpp
@@ -663,33 +663,38 @@ void CranedKeeper::PutNodeIntoUnavailList(const std::string &crane_id) {
 }
 
 void CranedKeeper::ConnectCranedNode_(CranedId const &craned_id) {
+  static Mutex s_craned_id_to_ip_cache_map_mtx;
   static std::unordered_map<CranedId, std::variant<ipv4_t, ipv6_t>>
       s_craned_id_to_ip_cache_map;
 
   std::string ip_addr;
 
-  auto it = s_craned_id_to_ip_cache_map.find(craned_id);
-  if (it != s_craned_id_to_ip_cache_map.end()) {
-    if (std::holds_alternative<ipv4_t>(it->second)) {  // Ipv4
-      ip_addr = crane::Ipv4ToStr(std::get<ipv4_t>(it->second));
-    } else {
-      CRANE_ASSERT(std::holds_alternative<ipv6_t>(it->second));
-      ip_addr = crane::Ipv6ToStr(std::get<ipv6_t>(it->second));
-    }
-  } else {
-    ipv4_t ipv4_addr;
-    ipv6_t ipv6_addr;
-    if (crane::ResolveIpv4FromHostname(craned_id, &ipv4_addr)) {
-      ip_addr = crane::Ipv4ToStr(ipv4_addr);
-      s_craned_id_to_ip_cache_map.emplace(craned_id, ipv4_addr);
-    } else if (crane::ResolveIpv6FromHostname(craned_id, &ipv6_addr)) {
-      ip_addr = crane::Ipv6ToStr(ipv6_addr);
-      s_craned_id_to_ip_cache_map.emplace(craned_id, ipv6_addr);
+  {
+    util::lock_guard guard(s_craned_id_to_ip_cache_map_mtx);
+
+    auto it = s_craned_id_to_ip_cache_map.find(craned_id);
+    if (it != s_craned_id_to_ip_cache_map.end()) {
+      if (std::holds_alternative<ipv4_t>(it->second)) {  // Ipv4
+        ip_addr = crane::Ipv4ToStr(std::get<ipv4_t>(it->second));
+      } else {
+        CRANE_ASSERT(std::holds_alternative<ipv6_t>(it->second));
+        ip_addr = crane::Ipv6ToStr(std::get<ipv6_t>(it->second));
+      }
     } else {
-      // Just hostname. It should never happen,
-      // but we add error handling here for robustness.
-      CRANE_ERROR("Unresolved hostname: {}", craned_id);
-      ip_addr = craned_id;
+      ipv4_t ipv4_addr;
+      ipv6_t ipv6_addr;
+      if (crane::ResolveIpv4FromHostname(craned_id, &ipv4_addr)) {
+        ip_addr = crane::Ipv4ToStr(ipv4_addr);
+        s_craned_id_to_ip_cache_map.emplace(craned_id, ipv4_addr);
+      } else if (crane::ResolveIpv6FromHostname(craned_id, &ipv6_addr)) {
+        ip_addr = crane::Ipv6ToStr(ipv6_addr);
+        s_craned_id_to_ip_cache_map.emplace(craned_id, ipv6_addr);
+      } else {
+        // Just hostname. It should never happen,
+        // but we add error handling here for robustness.
+        CRANE_ERROR("Unresolved hostname: {}", craned_id);
+        ip_addr = craned_id;
+      }
     }
   }
 
diff --git a/src/CraneCtld/TaskScheduler.h b/src/CraneCtld/TaskScheduler.h
index ec9bdfee3..32eab6508 100644
--- a/src/CraneCtld/TaskScheduler.h
+++ b/src/CraneCtld/TaskScheduler.h
@@ -263,6 +263,7 @@ class TaskScheduler {
       m_cancel_task_queue_.enqueue(
           CancelPendingTaskQueueElem{.task = std::move(pd_it->second)});
       m_cancel_task_async_handle_->send();
+      m_pending_task_map_.erase(pd_it);
       return CraneErr::kOk;
     }