SR-12833: processInfo.activeProcessorCount ignores CFS quotas

Sources/Foundation/ProcessInfo.swift

@@ -250,10 +250,17 @@ open class ProcessInfo: NSObject {
 
     internal let _processorCount: Int = Int(__CFProcessorCount())
     open var processorCount: Int { _processorCount }
-
+
+#if os(Linux)
+    // coreCount takes into account cgroup information eg if running under Docker
+    // __CFActiveProcessorCount uses sched_getaffinity() and sysconf(_SC_NPROCESSORS_ONLN)
+    internal let _activeProcessorCount: Int = ProcessInfo.coreCount() ?? Int(__CFActiveProcessorCount())
+#else
     internal let _activeProcessorCount: Int = Int(__CFActiveProcessorCount())
+#endif
+
     open var activeProcessorCount: Int { _activeProcessorCount }
-    
+
     internal let _physicalMemory = __CFMemorySize()
     open var physicalMemory: UInt64 {
         return _physicalMemory
@@ -293,6 +300,63 @@ open class ProcessInfo: NSObject {
     open var fullUserName: String {
         return NSFullUserName()
     }
+
+
+#if os(Linux)
+    // Support for CFS quotas for cpu count as used by Docker.
+    // Based on swift-nio code, https://github.com/apple/swift-nio/pull/1518
+    private static let cfsQuotaPath = "/sys/fs/cgroup/cpu/cpu.cfs_quota_us"
+    private static let cfsPeriodPath = "/sys/fs/cgroup/cpu/cpu.cfs_period_us"
+    private static let cpuSetPath = "/sys/fs/cgroup/cpuset/cpuset.cpus"
+
+    private static func firstLineOfFile(path: String) throws -> Substring {
+        let data = try Data(contentsOf: URL(fileURLWithPath: path))
+        if let string = String(data: data, encoding: .utf8), let line = string.split(separator: "\n").first {
+            return line
+        } else {
+            return ""
+        }
+    }
+
+    // These are internal access for testing
+    static func countCoreIds(cores: Substring) -> Int {
+        let ids = cores.split(separator: "-", maxSplits: 1)
+        guard let first = ids.first.flatMap({ Int($0, radix: 10) }),
+              let last = ids.last.flatMap({ Int($0, radix: 10) }),
+              last >= first
+        else { preconditionFailure("cpuset format is incorrect") }
+        return 1 + last - first
+    }
+
+    static func coreCount(cpuset cpusetPath: String) -> Int? {
+        guard let cpuset = try? firstLineOfFile(path: cpusetPath).split(separator: ","),
+              !cpuset.isEmpty
+        else { return nil }
+
+        return cpuset.map(countCoreIds).reduce(0, +)
+    }
+
+    static func coreCount(quota quotaPath: String,  period periodPath: String) -> Int? {
+        guard let quota = try? Int(firstLineOfFile(path: quotaPath)),
+              quota > 0
+        else { return nil }
+        guard let period = try? Int(firstLineOfFile(path: periodPath)),
+              period > 0
+        else { return nil }
+
+        return (quota - 1 + period) / period // always round up if fractional CPU quota requested
+    }
+
+    private static func coreCount() -> Int? {
+        if let quota = coreCount(quota: cfsQuotaPath, period: cfsPeriodPath) {
+            return quota
+        } else if let cpusetCount = coreCount(cpuset: cpuSetPath) {
+            return cpusetCount
+        } else {
+            return nil
+        }
+    }
+#endif
 }
 
 // SPI for TestFoundation

Tests/Foundation/Tests/TestProcessInfo.swift

@@ -7,17 +7,16 @@
 // See http://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
 //
 
+#if NS_FOUNDATION_ALLOWS_TESTABLE_IMPORT
+    #if canImport(SwiftFoundation) && !DEPLOYMENT_RUNTIME_OBJC
+        @testable import SwiftFoundation
+    #else
+        @testable import Foundation
+    #endif
+#endif
+
 class TestProcessInfo : XCTestCase {
 
-    static var allTests: [(String, (TestProcessInfo) -> () throws -> Void)] {
-        return [
-            ("test_operatingSystemVersion", test_operatingSystemVersion ),
-            ("test_processName", test_processName ),
-            ("test_globallyUniqueString", test_globallyUniqueString ),
-            ("test_environment", test_environment),
-        ]
-    }
-
     func test_operatingSystemVersion() {
         let processInfo = ProcessInfo.processInfo
         let versionString = processInfo.operatingSystemVersionString
@@ -134,4 +133,51 @@ class TestProcessInfo : XCTestCase {
         XCTAssertEqual(env["var4"], "x=")
         XCTAssertEqual(env["var5"], "=x=")
     }
+
+
+#if NS_FOUNDATION_ALLOWS_TESTABLE_IMPORT && os(Linux)
+    func test_cfquota_parsing() throws {
+
+        let tests = [
+            ("50000", "100000", 1),
+            ("100000", "100000", 1),
+            ("100000\n", "100000", 1),
+            ("100000", "100000\n", 1),
+            ("150000", "100000", 2),
+            ("200000", "100000", 2),
+            ("-1", "100000", nil),
+            ("100000", "-1", nil),
+            ("", "100000", nil),
+            ("100000", "", nil),
+            ("100000", "0", nil)
+        ]
+
+        try withTemporaryDirectory() { (_, tempDirPath) -> Void in
+            try tests.forEach { quota, period, count in
+                let (fd1, quotaPath) = try _NSCreateTemporaryFile(tempDirPath + "/quota")
+                FileHandle(fileDescriptor: fd1, closeOnDealloc: true).write(quota)
+
+                let (fd2, periodPath) = try _NSCreateTemporaryFile(tempDirPath + "/period")
+                FileHandle(fileDescriptor: fd2, closeOnDealloc: true).write(period)
+                XCTAssertEqual(ProcessInfo.coreCount(quota: quotaPath, period: periodPath), count)
+            }
+        }
+    }
+#endif
+
+
+    static var allTests: [(String, (TestProcessInfo) -> () throws -> Void)] {
+        var tests: [(String, (TestProcessInfo) -> () throws -> ())] = [
+            ("test_operatingSystemVersion", test_operatingSystemVersion ),
+            ("test_processName", test_processName ),
+            ("test_globallyUniqueString", test_globallyUniqueString ),
+            ("test_environment", test_environment),
+        ]
+
+#if NS_FOUNDATION_ALLOWS_TESTABLE_IMPORT && os(Linux)
+        tests.append(contentsOf: [ ("test_cfquota_parsing", test_cfquota_parsing) ])
+#endif
+
+        return tests
+    }
 }

Tests/Foundation/Utilities.swift

@@ -667,8 +667,10 @@ public func withTemporaryDirectory<R>(functionName: String = #function, block: (
         throw TestError.unexpectedNil
     }
 
-    let fname = String(functionName[..<idx])
-    let tmpDir = URL(fileURLWithPath: NSTemporaryDirectory(), isDirectory: true).appendingPathComponent(testBundleName()).appendingPathComponent(fname).appendingPathComponent(NSUUID().uuidString)
+    // Create the temporary directory as one level so that it doesnt leave a directory hierarchy on the filesystem
+    // eg tmp dir will be something like:  /tmp/TestFoundation-test_name-BE16B2FF-37FA-4F70-8A84-923D1CC2A860
+    let fname = testBundleName() + "-" + String(functionName[..<idx]) + "-" + NSUUID().uuidString
+    let tmpDir = URL(fileURLWithPath: NSTemporaryDirectory(), isDirectory: true).appendingPathComponent(fname)
     let fm = FileManager.default
     try? fm.removeItem(at: tmpDir)
     try fm.createDirectory(at: tmpDir, withIntermediateDirectories: true)