Bugfix: fix asan reporting use-after-poison

Signed-off-by: Li Junlin <[email protected]>

src/Craned/CgroupManager.cpp

@@ -383,40 +383,30 @@ bool CgroupManager::ReleaseCgroupByTaskIdOnly(task_id_t task_id) {
 }
 
 bool CgroupManager::ReleaseCgroup(uint32_t task_id, uid_t uid) {
-  if (!this->m_uid_to_task_ids_map_.Contains(uid)) {
-    CRANE_DEBUG(
-        "Trying to release a non-existent cgroup for uid #{}. Ignoring it...",
-        uid);
-    return false;
-  }
-
   this->m_task_id_to_cg_spec_map_.Erase(task_id);
 
   {
-    auto uid_task_id_map = this->m_uid_to_task_ids_map_.GetMapExclusivePtr();
-    auto task_id_set_ptr = uid_task_id_map->at(uid).GetExclusivePtr();
-
-    task_id_set_ptr->erase(task_id);
-    if (task_id_set_ptr->empty()) uid_task_id_map->erase(uid);
-  }
-
-  if (!this->m_task_id_to_cg_map_.Contains(task_id)) {
-    CRANE_DEBUG(
-        "Trying to release a non-existent cgroup for task #{}. Ignoring "
-        "it...",
-        task_id);
-
-    return false;
-  } else {
     // The termination of all processes in a cgroup is a time-consuming work.
     // Therefore, once we are sure that the cgroup for this task exists, we
     // let gRPC call return and put the termination work into the thread pool
     // to avoid blocking the event loop of TaskManager.
     // Kind of async behavior.
 
     // avoid deadlock by Erase at next line
-    Cgroup *cgroup = this->m_task_id_to_cg_map_[task_id]->release();
-    this->m_task_id_to_cg_map_.Erase(task_id);
+    auto task_id_to_cg_map_ptr =
+        this->m_task_id_to_cg_map_.GetMapExclusivePtr();
+    auto it = task_id_to_cg_map_ptr->find(task_id);
+    if (it == task_id_to_cg_map_ptr->end()) {
+      CRANE_DEBUG(
+          "Trying to release a non-existent cgroup for task #{}. Ignoring "
+          "it...",
+          task_id);
+
+      return false;
+    }
+    Cgroup *cgroup = it->second.GetExclusivePtr()->release();
+
+    task_id_to_cg_map_ptr->erase(task_id);
 
     if (cgroup != nullptr) {
       g_thread_pool->detach_task([cgroup]() {
@@ -442,8 +432,23 @@ bool CgroupManager::ReleaseCgroup(uint32_t task_id, uid_t uid) {
         delete cgroup;
       });
     }
-    return true;
   }
+
+  auto uid_task_ids_map_ptr = this->m_uid_to_task_ids_map_.GetMapExclusivePtr();
+  auto it = uid_task_ids_map_ptr->find(uid);
+  if (it == uid_task_ids_map_ptr->end()) {
+    CRANE_DEBUG(
+        "Trying to release a non-existent cgroup for uid #{}. Ignoring it...",
+        uid);
+    return false;
+  }
+
+  auto task_id_set_ptr = it->second.GetExclusivePtr();
+
+  task_id_set_ptr->erase(task_id);
+  if (task_id_set_ptr->empty()) uid_task_ids_map_ptr->erase(uid);
+
+  return true;
 }
 
 void CgroupManager::RmAllTaskCgroupsUnderController_(
@@ -480,8 +485,11 @@ bool CgroupManager::QueryTaskInfoOfUidAsync(uid_t uid, TaskInfoOfUid *info) {
   info->job_cnt = 0;
   info->cgroup_exists = false;
 
-  if (this->m_uid_to_task_ids_map_.Contains(uid)) {
-    auto task_ids = this->m_uid_to_task_ids_map_[uid];
+  if (auto task_ids = this->m_uid_to_task_ids_map_[uid]) {
+    if (!task_ids) {
+      CRANE_WARN("Uid {} not found in uid_to_task_ids_map", uid);
+      return false;
+    }
     info->job_cnt = task_ids->size();
     info->first_task_id = *task_ids->begin();
   }

src/Craned/CranedServer.cpp

@@ -174,11 +174,11 @@ grpc::Status CranedServiceImpl::CreateCgroupForTasks(
   for (int i = 0; i < request->task_id_list_size(); i++) {
     task_id_t task_id = request->task_id_list(i);
     uid_t uid = request->uid_list(i);
-    crane::grpc::ResourceInNode const &res = request->res_list(i);
+    crane::grpc::ResourceInNode res = request->res_list(i);
 
     CgroupSpec spec{.uid = uid,
                     .task_id = task_id,
-                    .res_in_node = std::move(res),
+                    .res_in_node = res,
                     .execution_node = request->execution_node(i)};
     CRANE_TRACE("Receive CreateCgroup for task #{}, uid {}", task_id, uid);
     cg_specs.emplace_back(std::move(spec));