Backport #66237 to 24.6: Fix detection of number of CPUs in containers

base/base/cgroupsv2.cpp

@@ -3,8 +3,9 @@
 #include <base/defines.h>
 
 #include <fstream>
-#include <sstream>
+#include <string>
 
+namespace fs = std::filesystem;
 
 bool cgroupsV2Enabled()
 {
@@ -13,11 +14,11 @@ bool cgroupsV2Enabled()
     {
         /// This file exists iff the host has cgroups v2 enabled.
         auto controllers_file = default_cgroups_mount / "cgroup.controllers";
-        if (!std::filesystem::exists(controllers_file))
+        if (!fs::exists(controllers_file))
             return false;
         return true;
     }
-    catch (const std::filesystem::filesystem_error &) /// all "underlying OS API errors", typically: permission denied
+    catch (const fs::filesystem_error &) /// all "underlying OS API errors", typically: permission denied
     {
         return false; /// not logging the exception as most callers fall back to cgroups v1
     }
@@ -33,8 +34,9 @@ bool cgroupsV2MemoryControllerEnabled()
     /// According to https://docs.kernel.org/admin-guide/cgroup-v2.html, file "cgroup.controllers" defines which controllers are available
     /// for the current + child cgroups. The set of available controllers can be restricted from level to level using file
     /// "cgroups.subtree_control". It is therefore sufficient to check the bottom-most nested "cgroup.controllers" file.
-    std::string cgroup = cgroupV2OfProcess();
-    auto cgroup_dir = cgroup.empty() ? default_cgroups_mount : (default_cgroups_mount / cgroup);
+    fs::path cgroup_dir = cgroupV2PathOfProcess();
+    if (cgroup_dir.empty())
+        return false;
     std::ifstream controllers_file(cgroup_dir / "cgroup.controllers");
     if (!controllers_file.is_open())
         return false;
@@ -46,7 +48,7 @@ bool cgroupsV2MemoryControllerEnabled()
 #endif
 }
 
-std::string cgroupV2OfProcess()
+fs::path cgroupV2PathOfProcess()
 {
 #if defined(OS_LINUX)
     chassert(cgroupsV2Enabled());
@@ -54,17 +56,18 @@ std::string cgroupV2OfProcess()
     /// A simpler way to get the membership is:
     std::ifstream cgroup_name_file("/proc/self/cgroup");
     if (!cgroup_name_file.is_open())
-        return "";
+        return {};
     /// With cgroups v2, there will be a *single* line with prefix "0::/"
     /// (see https://docs.kernel.org/admin-guide/cgroup-v2.html)
     std::string cgroup;
     std::getline(cgroup_name_file, cgroup);
     static const std::string v2_prefix = "0::/";
     if (!cgroup.starts_with(v2_prefix))
-        return "";
+        return {};
     cgroup = cgroup.substr(v2_prefix.length());
-    return cgroup;
+    /// Note: The 'root' cgroup can have an empty cgroup name, this is valid
+    return default_cgroups_mount / cgroup;
 #else
-    return "";
+    return {};
 #endif
 }

base/base/cgroupsv2.h

@@ -1,7 +1,6 @@
 #pragma once
 
 #include <filesystem>
-#include <string>
 
 #if defined(OS_LINUX)
 /// I think it is possible to mount the cgroups hierarchy somewhere else (e.g. when in containers).
@@ -16,7 +15,7 @@ bool cgroupsV2Enabled();
 /// Assumes that cgroupsV2Enabled() is enabled.
 bool cgroupsV2MemoryControllerEnabled();
 
-/// Which cgroup does the process belong to?
-/// Returns an empty string if the cgroup cannot be determined.
+/// Detects which cgroup v2 the process belongs to and returns the filesystem path to the cgroup.
+/// Returns an empty path the cgroup cannot be determined.
 /// Assumes that cgroupsV2Enabled() is enabled.
-std::string cgroupV2OfProcess();
+std::filesystem::path cgroupV2PathOfProcess();

base/base/getMemoryAmount.cpp

@@ -23,8 +23,9 @@ std::optional<uint64_t> getCgroupsV2MemoryLimit()
     if (!cgroupsV2MemoryControllerEnabled())
         return {};
 
-    std::string cgroup = cgroupV2OfProcess();
-    auto current_cgroup = cgroup.empty() ? default_cgroups_mount : (default_cgroups_mount / cgroup);
+    std::filesystem::path current_cgroup = cgroupV2PathOfProcess();
+    if (current_cgroup.empty())
+        return {};
 
     /// Open the bottom-most nested memory limit setting file. If there is no such file at the current
     /// level, try again at the parent level as memory settings are inherited.

src/Common/CgroupsMemoryUsageObserver.cpp

@@ -22,6 +22,7 @@
 #define STRINGIFY(x) STRINGIFY_HELPER(x)
 #endif
 
+namespace fs = std::filesystem;
 
 namespace DB
 {
@@ -125,15 +126,16 @@ std::optional<std::string> getCgroupsV2FileName()
     if (!cgroupsV2MemoryControllerEnabled())
         return {};
 
-    String cgroup = cgroupV2OfProcess();
-    auto current_cgroup = cgroup.empty() ? default_cgroups_mount : (default_cgroups_mount / cgroup);
+    fs::path current_cgroup = cgroupV2PathOfProcess();
+    if (current_cgroup.empty())
+        return {};
 
     /// Return the bottom-most nested current memory file. If there is no such file at the current
     /// level, try again at the parent level as memory settings are inherited.
     while (current_cgroup != default_cgroups_mount.parent_path())
     {
         auto path = current_cgroup / "memory.current";
-        if (std::filesystem::exists(path))
+        if (fs::exists(path))
             return {path};
         current_cgroup = current_cgroup.parent_path();
     }
@@ -143,7 +145,7 @@ std::optional<std::string> getCgroupsV2FileName()
 std::optional<std::string> getCgroupsV1FileName()
 {
     auto path = default_cgroups_mount / "memory/memory.stat";
-    if (!std::filesystem::exists(path))
+    if (!fs::exists(path))
         return {};
     return {path};
 }

src/Common/getNumberOfPhysicalCPUCores.cpp

@@ -37,12 +37,12 @@ uint32_t getCGroupLimitedCPUCores(unsigned default_cpu_count)
     /// cgroupsv2
     if (cgroupsV2Enabled())
     {
-        /// First, we identify the cgroup the process belongs
-        std::string cgroup = cgroupV2OfProcess();
-        if (cgroup.empty())
+        /// First, we identify the path of the cgroup the process belongs
+        std::filesystem::path cgroup_path = cgroupV2PathOfProcess();
+        if (cgroup_path.empty())
             return default_cpu_count;
 
-        auto current_cgroup = cgroup.empty() ? default_cgroups_mount : (default_cgroups_mount / cgroup);
+        auto current_cgroup = cgroup_path;
 
         // Looking for cpu.max in directories from the current cgroup to the top level
         // It does not stop on the first time since the child could have a greater value than parent
@@ -62,7 +62,7 @@ uint32_t getCGroupLimitedCPUCores(unsigned default_cpu_count)
             }
             current_cgroup = current_cgroup.parent_path();
         }
-        current_cgroup = default_cgroups_mount / cgroup;
+        current_cgroup = cgroup_path;
         // Looking for cpuset.cpus.effective in directories from the current cgroup to the top level
         while (current_cgroup != default_cgroups_mount.parent_path())
         {