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

Sources/CNIOLinux/include/CNIOLinux.h

@@ -25,6 +25,8 @@
 #include <sys/xattr.h>
 #include <sys/types.h>
 #include <sys/stat.h>
+#include <sys/vfs.h>
+#include <linux/magic.h>
 #include <sched.h>
 #include <stdbool.h>
 #include <errno.h>
@@ -149,6 +151,10 @@ extern const unsigned long CNIOLinux_UTIME_NOW;
 
 extern const long CNIOLinux_UDP_MAX_SEGMENTS;
 
+// Filesystem magic constants for cgroup detection
+extern const unsigned long CNIOLinux_TMPFS_MAGIC;
+extern const unsigned long CNIOLinux_CGROUP2_SUPER_MAGIC;
+
 // 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,8 @@ 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;
 
+const unsigned long CNIOLinux_TMPFS_MAGIC = TMPFS_MAGIC;
+const unsigned long CNIOLinux_CGROUP2_SUPER_MAGIC = CGROUP2_SUPER_MAGIC;
 
 #ifdef UDP_MAX_SEGMENTS
 const long CNIOLinux_UDP_MAX_SEGMENTS = UDP_MAX_SEGMENTS;

Sources/NIOCore/Linux.swift

@@ -20,12 +20,79 @@ import CNIOLinux
 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)
+    static let cpuSetPathV1 = "/sys/fs/cgroup/cpuset/cpuset.cpus"
+    static lazy let cpuSetPathV2: String? = {
+        if let cgroupV2MountPoint = Self.cgroupV2MountPoint {
+            return NIOFilePath(cgroupV2MountPoint).appending(["cpuset.cpus"]).description
+        }
+        return nil
+    }()
+
+    static lazy let cgroupV2MountPoint: String? = {
+        guard let fh = try? NIOFileHandle(_deprecatedPath: "/proc/self/cgroup") else { return nil }
         defer { try! fh.close() }
+        guard let lines = try? Self.readLines(fh: fh) else { return nil }
+
+        // Parse each line looking for cgroup v2 format: "0::/path"
+        for line in lines {
+            if let cgroupPath = Self.parseV2CgroupLine(line) {
+                return cgroupPath
+            }
+        }
+
+        return nil
+    }()
+
+    /// Returns the appropriate cpuset path based on the detected cgroup version
+    static lazy 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 func 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 Int(CNIOLinux.TMPFS_MAGIC):
+            return .v1
+        case Int(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
+    private static func parseV2CgroupLine(_ line: Substring) -> String? {
+        // Expected format is "0::/path"
+        let parts = line.split(separator: ":", maxSplits: 2)
+
+        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(fh: NIOFileHandle) 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
@@ -39,7 +106,13 @@ enum Linux {
                 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? {
+        let fh = try NIOFileHandle(_deprecatedPath: path)
+        defer { try! fh.close() }
+        return try? Self.readLines(fh: fh).first
     }
 
     private static func countCoreIds(cores: Substring) -> Int {
@@ -54,7 +127,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 +140,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 {
+        return try syscall(blocking: false) {
+            sysStatfs(path, &buf)
+        }.result
+    }
+    #endif
+
     #endif  // !os(WASI)
 }

Sources/NIOCore/Utilities.swift

@@ -121,11 +121,25 @@ 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))

Tests/NIOCoreTests/LinuxTest.swift

@@ -2,7 +2,7 @@
 //
 // This source file is part of the SwiftNIO open source project
 //
-// Copyright (c) 2017-2023 Apple Inc. and the SwiftNIO project authors
+// Copyright (c) 2017-2025 Apple Inc. and the SwiftNIO project authors
 // Licensed under Apache License v2.0
 //
 // See LICENSE.txt for license information
@@ -17,6 +17,113 @@ import XCTest
 @testable import NIOCore
 
 class LinuxTest: XCTestCase {
+    func testCoreCountCgroup1RestrictionWithVariousIntegerFormats() throws {
+        #if os(Linux) || os(Android)
+        // Test various integer formats in cgroup files
+        let testCases = [
+            ("42", "100000", 1),      // Simple integer
+            ("0", "100000", nil),     // Zero should be rejected
+        ]
+
+        for (quotaContent, periodContent, expectedResult) in testCases {
+            try withTemporaryFile(content: quotaContent) { (_, quotaPath) -> Void in
+                try withTemporaryFile(content: periodContent) { (_, periodPath) -> Void in
+                    let result = Linux.coreCountCgroup1Restriction(quota: quotaPath, period: periodPath)
+                    XCTAssertEqual(result, expectedResult, "Failed for quota '\(quotaContent)'")
+                }
+            }
+        }
+
+        // Test invalid integer cases
+        let invalidCases = ["abc", "12abc", ""]
+        for invalidContent in invalidCases {
+            try withTemporaryFile(content: invalidContent) { (_, quotaPath) -> Void in
+                try withTemporaryFile(content: "100000") { (_, periodPath) -> Void in
+                    let result = Linux.coreCountCgroup1Restriction(quota: quotaPath, period: periodPath)
+                    XCTAssertNil(result, "Should return nil for invalid quota '\(invalidContent)'")
+                }
+            }
+        }
+        #endif
+    }
+
+    func testCoreCountWithMultipleRanges() throws {
+        #if os(Linux) || os(Android)
+        // Test coreCount function with multiple CPU ranges
+        let content = "0,2,4-6"  // Should count as 5 cores: 0,2,4,5,6
+        try withTemporaryFile(content: content) { (_, path) -> Void in
+            let result = Linux.coreCount(cpuset: path)
+            XCTAssertEqual(result, 5)
+        }
+        #endif
+    }
+
+    func testCoreCountWithSingleRange() throws {
+        #if os(Linux) || os(Android)
+        let content = "0-3"  // Should count as 4 cores
+        try withTemporaryFile(content: content) { (_, path) -> Void in
+            let result = Linux.coreCount(cpuset: path)
+            XCTAssertEqual(result, 4)
+        }
+        #endif
+    }
+
+    func testCoreCountWithEmptyFile() throws {
+        #if os(Linux) || os(Android)
+        try withTemporaryFile(content: "") { (_, path) -> Void in
+            let result = Linux.coreCount(cpuset: path)
+            XCTAssertNil(result)  // Empty file should return nil
+        }
+        #endif
+    }
+
+    func testCoreCountReadsOnlyFirstLine() throws {
+        #if os(Linux) || os(Android)
+        // Test that coreCount only processes the first line of the file
+        let content = "0-1\n2-3\n4-5"  // First line should be "0-1" = 2 cores
+        try withTemporaryFile(content: content) { (_, path) -> Void in
+            let result = Linux.coreCount(cpuset: path)
+            XCTAssertEqual(result, 2)  // Should only process first line
+        }
+        #endif
+    }
+
+    func testCoreCountWithSimpleCpuset() throws {
+        #if os(Linux) || os(Android)
+        // Test coreCount function with simple cpuset formats
+        let testCases = [
+            ("0", 1),      // Single core
+            ("0-3", 4),    // Range 0,1,2,3
+            ("5-7", 3),    // Range 5,6,7
+            ("10-10", 1),  // Single core as range
+        ]
+
+        for (input, expected) in testCases {
+            try withTemporaryFile(content: input) { (_, path) -> Void in
+                let result = Linux.coreCount(cpuset: path)
+                XCTAssertEqual(result, expected, "Failed for input '\(input)'")
+            }
+        }
+        #endif
+    }
+
+    func testCoreCountWithComplexCpuset() throws {
+        #if os(Linux) || os(Android)
+        // Test more complex cpuset formats
+        let cpusets = [
+            ("0", 1),
+            ("0,2", 2),
+            ("0-1,4-5", 4),  // Two ranges: 0,1 and 4,5
+            ("0,2-4,7,9-10", 7),  // Mixed: 0, 2,3,4, 7, 9,10
+        ]
+        for (cpuset, count) in cpusets {
+            try withTemporaryFile(content: cpuset) { (_, path) -> Void in
+                XCTAssertEqual(Linux.coreCount(cpuset: path), count)
+            }
+        }
+        #endif
+    }
+
     func testCoreCountQuota() throws {
         #if os(Linux) || os(Android)
         let coreCountQuoats = [
@@ -61,7 +168,7 @@ class LinuxTest: XCTestCase {
         #endif
     }
 
-    func testCoreCountCgoup2() throws {
+    func testCoreCountCgroup2() throws {
         #if os(Linux) || os(Android)
         let contents = [
             ("max 100000", nil),