Fix coreCount on Linux when using cgroup v2 with CFS throttling disabled

Sources/CNIOLinux/include/CNIOLinux.h

@@ -25,11 +25,15 @@
 #include <sys/xattr.h>
 #include <sys/types.h>
 #include <sys/stat.h>
+#include <sys/vfs.h>
 #include <sched.h>
 #include <stdbool.h>
 #include <errno.h>
 #include <pthread.h>
 #include <netinet/ip.h>
+#if __has_include(<linux/magic.h>)
+#include <linux/magic.h>
+#endif
 #if __has_include(<linux/udp.h>)
 #include <linux/udp.h>
 #else
@@ -149,6 +153,25 @@ extern const unsigned long CNIOLinux_UTIME_NOW;
 
 extern const long CNIOLinux_UDP_MAX_SEGMENTS;
 
+// Filesystem magic constants for cgroup detection
+#ifdef __ANDROID__
+#if defined(__LP64__)
+extern const uint64_t CNIOLinux_TMPFS_MAGIC;
+extern const uint64_t CNIOLinux_CGROUP2_SUPER_MAGIC;
+#else
+extern const uint32_t CNIOLinux_TMPFS_MAGIC;
+extern const uint32_t CNIOLinux_CGROUP2_SUPER_MAGIC;
+#endif
+#else
+#ifdef __FSWORD_T_TYPE
+extern const __fsword_t CNIOLinux_TMPFS_MAGIC;
+extern const __fsword_t CNIOLinux_CGROUP2_SUPER_MAGIC;
+#else
+extern const unsigned long CNIOLinux_TMPFS_MAGIC;
+extern const unsigned long CNIOLinux_CGROUP2_SUPER_MAGIC;
+#endif
+#endif
+
 // A workaround for incorrect nullability annotations in the Android SDK.
 FTS *CNIOLinux_fts_open(char * const *path_argv, int options, int (*compar)(const FTSENT **, const FTSENT **));
 

Sources/CNIOLinux/shim.c

@@ -215,6 +215,30 @@ const int CNIOLinux_AT_EMPTY_PATH = AT_EMPTY_PATH;
 const unsigned long CNIOLinux_UTIME_OMIT = UTIME_OMIT;
 const unsigned long CNIOLinux_UTIME_NOW = UTIME_NOW;
 
+#ifndef TMPFS_MAGIC
+#define TMPFS_MAGIC 0x01021994
+#endif
+#ifndef CGROUP2_SUPER_MAGIC
+#define CGROUP2_SUPER_MAGIC 0x63677270
+#endif
+
+#ifdef __ANDROID__
+#if defined(__LP64__)
+const uint64_t CNIOLinux_TMPFS_MAGIC = TMPFS_MAGIC;
+const uint64_t CNIOLinux_CGROUP2_SUPER_MAGIC = CGROUP2_SUPER_MAGIC;
+#else
+const uint32_t CNIOLinux_TMPFS_MAGIC = TMPFS_MAGIC;
+const uint32_t CNIOLinux_CGROUP2_SUPER_MAGIC = CGROUP2_SUPER_MAGIC;
+#endif
+#else
+#ifdef __FSWORD_T_TYPE
+const __fsword_t CNIOLinux_TMPFS_MAGIC = TMPFS_MAGIC;
+const __fsword_t CNIOLinux_CGROUP2_SUPER_MAGIC = CGROUP2_SUPER_MAGIC;
+#else
+const unsigned long CNIOLinux_TMPFS_MAGIC = TMPFS_MAGIC;
+const unsigned long CNIOLinux_CGROUP2_SUPER_MAGIC = CGROUP2_SUPER_MAGIC;
+#endif
+#endif
 
 #ifdef UDP_MAX_SEGMENTS
 const long CNIOLinux_UDP_MAX_SEGMENTS = UDP_MAX_SEGMENTS;

Sources/NIOCore/Linux.swift

@@ -17,29 +17,111 @@
 
 #if os(Linux) || os(Android)
 import CNIOLinux
+
+#if canImport(Android)
+@preconcurrency import Android
+#endif
+
 enum Linux {
     static let cfsQuotaPath = "/sys/fs/cgroup/cpu/cpu.cfs_quota_us"
     static let cfsPeriodPath = "/sys/fs/cgroup/cpu/cpu.cfs_period_us"
-    static let cpuSetPath = "/sys/fs/cgroup/cpuset/cpuset.cpus"
     static let cfsCpuMaxPath = "/sys/fs/cgroup/cpu.max"
 
-    private static func firstLineOfFile(path: String) throws -> Substring {
-        let fh = try NIOFileHandle(_deprecatedPath: path)
-        defer { try! fh.close() }
+    static let cpuSetPathV1 = "/sys/fs/cgroup/cpuset/cpuset.cpus"
+    static let cpuSetPathV2: String? = {
+        if let cgroupV2MountPoint = Self.cgroupV2MountPoint {
+            return "\(cgroupV2MountPoint)/cpuset.cpus"
+        }
+        return nil
+    }()
+
+    static let cgroupV2MountPoint: String? = {
+        guard
+            let fd = try? SystemCalls.open(file: "/proc/self/cgroup", oFlag: O_RDONLY, mode: NIOPOSIXFileMode(S_IRUSR))
+        else { return nil }
+        defer { try! SystemCalls.close(descriptor: fd) }
+        guard let lines = try? Self.readLines(descriptor: fd) else { return nil }
+
+        // Parse each line looking for cgroup v2 format: "0::/path"
+        for line in lines {
+            if let cgroupPath = Self.parseV2CgroupLine(line) {
+                return "/sys/fs/cgroup\(cgroupPath)"
+            }
+        }
+
+        return nil
+    }()
+
+    /// Returns the appropriate cpuset path based on the detected cgroup version
+    static let cpuSetPath: String? = {
+        guard let version = Self.cgroupVersion else { return nil }
+
+        switch version {
+        case .v1:
+            return cpuSetPathV1
+        case .v2:
+            return cpuSetPathV2
+        }
+    }()
+
+    /// Detects whether we're using cgroup v1 or v2
+    static let cgroupVersion: CgroupVersion? = {
+        var fs = statfs()
+        guard let result = try? SystemCalls.statfs("/sys/fs/cgroup", &fs), result == 0 else { return nil }
+
+        switch fs.f_type {
+        case CNIOLinux_TMPFS_MAGIC:
+            return .v1
+        case CNIOLinux_CGROUP2_SUPER_MAGIC:
+            return .v2
+        default:
+            return nil
+        }
+    }()
+
+    enum CgroupVersion {
+        case v1
+        case v2
+    }
+
+    /// Parses a single line from /proc/self/cgroup to extract cgroup v2 path
+    internal static func parseV2CgroupLine(_ line: Substring) -> String? {
+        // Expected format is "0::/path"
+        let parts = line.split(separator: ":", maxSplits: 2, omittingEmptySubsequences: false)
+
+        guard parts.count == 3,
+            parts[0] == "0",
+            parts[1] == ""
+        else {
+            return nil
+        }
+
+        // Extract the path from parts[2]
+        return String(parts[2])
+    }
+
+    private static func readLines(descriptor: CInt) throws -> [Substring] {
         // linux doesn't properly report /sys/fs/cgroup/* files lengths so we use a reasonable limit
         var buf = ByteBufferAllocator().buffer(capacity: 1024)
         try buf.writeWithUnsafeMutableBytes(minimumWritableBytes: buf.capacity) { ptr in
-            let res = try fh.withUnsafeFileDescriptor { fd -> CoreIOResult<ssize_t> in
-                try SystemCalls.read(descriptor: fd, pointer: ptr.baseAddress!, size: ptr.count)
-            }
+            let res = try SystemCalls.read(descriptor: descriptor, pointer: ptr.baseAddress!, size: ptr.count)
+
             switch res {
             case .processed(let n):
                 return n
             case .wouldBlock:
                 preconditionFailure("read returned EWOULDBLOCK despite a blocking fd")
             }
         }
-        return String(buffer: buf).prefix(while: { $0 != "\n" })
+        return String(buffer: buf).split(separator: "\n")
+    }
+
+    private static func firstLineOfFile(path: String) throws -> Substring? {
+        guard let fd = try? SystemCalls.open(file: path, oFlag: O_RDONLY, mode: NIOPOSIXFileMode(S_IRUSR)) else {
+            return nil
+        }
+        defer { try! SystemCalls.close(descriptor: fd) }
+        return try? Self.readLines(descriptor: fd).first
     }
 
     private static func countCoreIds(cores: Substring) -> Int {
@@ -54,7 +136,7 @@ enum Linux {
 
     static func coreCount(cpuset cpusetPath: String) -> Int? {
         guard
-            let cpuset = try? firstLineOfFile(path: cpusetPath).split(separator: ","),
+            let cpuset = try? firstLineOfFile(path: cpusetPath).flatMap({ $0.split(separator: ",") }),
             !cpuset.isEmpty
         else { return nil }
         return cpuset.map(countCoreIds).reduce(0, +)
@@ -67,11 +149,11 @@ enum Linux {
         period periodPath: String = Linux.cfsPeriodPath
     ) -> Int? {
         guard
-            let quota = try? Int(firstLineOfFile(path: quotaPath)),
+            let quota = try? firstLineOfFile(path: quotaPath).flatMap({ Int($0) }),
             quota > 0
         else { return nil }
         guard
-            let period = try? Int(firstLineOfFile(path: periodPath)),
+            let period = try? firstLineOfFile(path: periodPath).flatMap({ Int($0) }),
             period > 0
         else { return nil }
         return (quota - 1 + period) / period  // always round up if fractional CPU quota requested

Sources/NIOCore/SystemCallHelpers.swift

@@ -35,6 +35,11 @@ import CNIOWindows
 #error("The system call helpers module was unable to identify your C library.")
 #endif
 
+#if os(Linux) || os(Android)
+import CNIOLinux
+private let sysStatfs: @convention(c) (UnsafePointer<CChar>, UnsafeMutablePointer<statfs>) -> CInt = statfs
+#endif
+
 #if os(Windows)
 private let sysDup: @convention(c) (CInt) -> CInt = _dup
 private let sysClose: @convention(c) (CInt) -> CInt = _close
@@ -232,5 +237,19 @@ enum SystemCalls {
         }
     }
     #endif
+
+    #if os(Linux) || os(Android)
+    @inline(never)
+    @usableFromInline
+    internal static func statfs(
+        _ path: UnsafePointer<CChar>,
+        _ buf: inout statfs
+    ) throws -> CInt {
+        try syscall(blocking: false) {
+            sysStatfs(path, &buf)
+        }.result
+    }
+    #endif
+
     #endif  // !os(WASI)
 }

Sources/NIOCore/Utilities.swift

@@ -121,11 +121,26 @@ public enum System: Sendable {
         .map { $0.ProcessorMask.nonzeroBitCount }
         .reduce(0, +)
         #elseif os(Linux) || os(Android)
-        if let quota2 = Linux.coreCountCgroup2Restriction() {
-            return quota2
-        } else if let quota = Linux.coreCountCgroup1Restriction() {
-            return quota
-        } else if let cpusetCount = Linux.coreCount(cpuset: Linux.cpuSetPath) {
+        var cpuSetPath: String?
+
+        switch Linux.cgroupVersion {
+        case .v1:
+            if let quota = Linux.coreCountCgroup1Restriction() {
+                return quota
+            }
+            cpuSetPath = Linux.cpuSetPathV1
+        case .v2:
+            if let quota = Linux.coreCountCgroup2Restriction() {
+                return quota
+            }
+            cpuSetPath = Linux.cpuSetPathV2
+        case .none:
+            break
+        }
+
+        if let cpuSetPath,
+            let cpusetCount = Linux.coreCount(cpuset: cpuSetPath)
+        {
             return cpusetCount
         } else {
             return sysconf(CInt(_SC_NPROCESSORS_ONLN))