refactor: Refactor cgroup controller check on craned start.

Signed-off-by: Li Junlin <xiafeng.li@foxmail.com>
PKUHPC · Dec 12, 2024 · c75c0bd · c75c0bd
1 parent 5a61084
commit c75c0bd
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 31 deletions.
diff --git a/src/Craned/CgroupManager.cpp b/src/Craned/CgroupManager.cpp
@@ -208,44 +208,58 @@ void CgroupManager::RmAllTaskCgroups_() {
       CgroupConstant::Controller::DEVICES_CONTROLLER);
 }
 
-void CgroupManager::ControllersMounted() {
+bool CgroupManager::ControllersMounted() {
   using namespace CgroupConstant;
+  auto all_controller_mounted = true;
   if (cg_version_ == CgroupVersion::CGROUP_V1) {
     if (!Mounted(Controller::BLOCK_CONTROLLER)) {
-      CRANE_WARN("Cgroup controller for I/O statistics is not available.");
+      CRANE_WARN("CgroupV1 controller for I/O statistics is not available.");
+      all_controller_mounted = false;
     }
     if (!Mounted(Controller::FREEZE_CONTROLLER)) {
-      CRANE_WARN("Cgroup controller for process management is not available.");
+      CRANE_WARN(
+          "CgroupV1 controller for process management is not available.");
+      all_controller_mounted = false;
     }
     if (!Mounted(Controller::CPUACCT_CONTROLLER)) {
-      CRANE_WARN("Cgroup controller for CPU accounting is not available.");
+      CRANE_WARN("CgroupV1 controller for CPU accounting is not available.");
+      all_controller_mounted = false;
     }
     if (!Mounted(Controller::MEMORY_CONTROLLER)) {
-      CRANE_WARN("Cgroup controller for memory accounting is not available.");
+      CRANE_WARN("CgroupV1 controller for memory accounting is not available.");
+      all_controller_mounted = false;
     }
     if (!Mounted(Controller::CPU_CONTROLLER)) {
-      CRANE_WARN("Cgroup controller for CPU is not available.");
+      CRANE_WARN("CgroupV1 controller for CPU is not available.");
+      all_controller_mounted = false;
     }
     if (!Mounted(Controller::DEVICES_CONTROLLER)) {
-      CRANE_WARN("Cgroup controller for DEVICES is not available.");
+      CRANE_WARN("CgroupV1 controller for DEVICES is not available.");
+      all_controller_mounted = false;
     }
   } else if (cg_version_ == CgroupVersion::CGROUP_V2) {
     if (!Mounted(Controller::CPU_CONTROLLER_V2)) {
-      CRANE_WARN("Cgroup controller for CPU is not available.");
+      CRANE_WARN("CgroupV2 controller for CPU is not available.");
+      all_controller_mounted = false;
     }
     if (!Mounted(Controller::MEMORY_CONTORLLER_V2)) {
-      CRANE_WARN("Cgroup controller for memory is not available.");
+      CRANE_WARN("CgroupV2 controller for memory is not available.");
+      all_controller_mounted = false;
     }
     if (!Mounted(Controller::CPUSET_CONTROLLER_V2)) {
-      CRANE_WARN("Cgroup controller for cpuset is not available.");
+      CRANE_WARN("CgroupV2 controller for cpuset is not available.");
+      all_controller_mounted = false;
     }
     if (!Mounted(Controller::IO_CONTROLLER_V2)) {
-      CRANE_WARN("Cgroup controller for I/O statistics is not available.");
+      CRANE_WARN("CgroupV2 controller for I/O statistics is not available.");
+      all_controller_mounted = false;
     }
     if (!Mounted(Controller::PIDS_CONTROLLER_V2)) {
-      CRANE_WARN("Cgroup controller for pids is not available.");
+      CRANE_WARN("CgroupV2 controller for pids is not available.");
+      all_controller_mounted = false;
     }
   }
+  return all_controller_mounted;
 }
 
 /*
@@ -501,9 +515,8 @@ bool CgroupManager::AllocateAndGetCgroup(task_id_t task_id,
   }
 
   if (g_config.Plugin.Enabled) {
-    g_plugin_client->CreateCgroupHookAsync(task_id, 
-                                          pcg->GetCgroupString(),
-                                          res.dedicated_res_in_node());
+    g_plugin_client->CreateCgroupHookAsync(task_id, pcg->GetCgroupString(),
+                                           res.dedicated_res_in_node());
   }
 
   CRANE_TRACE(
@@ -596,7 +609,8 @@ bool CgroupManager::ReleaseCgroup(uint32_t task_id, uid_t uid) {
     CgroupInterface *cgroup = it->second.GetExclusivePtr()->release();
 
     if (g_config.Plugin.Enabled) {
-      g_plugin_client->DestroyCgroupHookAsync(task_id, cgroup->GetCgroupString());
+      g_plugin_client->DestroyCgroupHookAsync(task_id,
+                                              cgroup->GetCgroupString());
     }
 
     task_id_to_cg_map_ptr->erase(task_id);
@@ -628,7 +642,8 @@ bool CgroupManager::ReleaseCgroup(uint32_t task_id, uid_t uid) {
   }
 
   {
-    auto uid_task_ids_map_ptr = this->m_uid_to_task_ids_map_.GetMapExclusivePtr();
+    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(

diff --git a/src/Craned/CgroupManager.h b/src/Craned/CgroupManager.h
@@ -465,7 +465,7 @@ class CgroupManager {
     return bool(m_mounted_controllers_ & ControllerFlags{controller});
   }
 
-  void ControllersMounted();
+  bool ControllersMounted();
 
   bool QueryTaskInfoOfUidAsync(uid_t uid, TaskInfoOfUid *info);
 

diff --git a/src/Craned/Craned.cpp b/src/Craned/Craned.cpp
@@ -603,19 +603,7 @@ void GlobalVariableInit() {
   using Craned::CgroupConstant::Controller;
   g_cg_mgr = std::make_unique<Craned::CgroupManager>();
   g_cg_mgr->Init();
-  if (g_cg_mgr->GetCgroupVersion() ==
-          Craned::CgroupConstant::CgroupVersion::CGROUP_V1 &&
-      (!g_cg_mgr->Mounted(Controller::CPU_CONTROLLER) ||
-       !g_cg_mgr->Mounted(Controller::MEMORY_CONTROLLER) ||
-       !g_cg_mgr->Mounted(Controller::DEVICES_CONTROLLER))) {
-    CRANE_ERROR("Failed to initialize cpu,memory,devices cgroups controller.");
-    std::exit(1);
-  }
-  if (g_cg_mgr->GetCgroupVersion() ==
-          Craned::CgroupConstant::CgroupVersion::CGROUP_V2 &&
-      (!g_cg_mgr->Mounted(Controller::CPU_CONTROLLER_V2) ||
-       !g_cg_mgr->Mounted(Controller::MEMORY_CONTORLLER_V2))) {
-    CRANE_ERROR("Failed to initialize cpu,memory cgroups controller.");
+  if (!g_cg_mgr->ControllersMounted()) {
     std::exit(1);
   }