fix cores count in containers (#1518)

Motivation:

The current System.coreCount implementation relies on _SC_NPROCESSORS_ONLN so isn't cgroups aware which might have a bad impact for apps runnings in containers (e.g. Docker, Kubernetes, Amazon ECS...).

Modifications:

- Changed System.coreCount on Linux only to make it read from CFS quotas and cpusets when present.
- Removed incorrect precondition in Sources/NIO/LinuxCPUSet.swift

Result:

System.coreCount returns correct values when apps run in containers.

Co-authored-by: Johannes Weiss <johannesweiss@apple.com>

Sources/NIO/Linux.swift

@@ -113,6 +113,9 @@ internal enum Epoll {
 }
 
 internal 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"
 #if os(Android)
     static let SOCK_CLOEXEC = Glibc.SOCK_CLOEXEC
     static let SOCK_NONBLOCK = Glibc.SOCK_NONBLOCK
@@ -132,5 +135,54 @@ internal enum Linux {
         }
         return fd
     }
+
+    private static func firstLineOfFile(path: String) throws -> Substring {
+        let fh = try NIOFileHandle(path: path)
+        defer { try! fh.close() }
+        // 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 -> IOResult<ssize_t> in
+                return try Posix.read(descriptor: fd, 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" })
+    }
+
+    private 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
+    }
 }
 #endif

Sources/NIO/LinuxCPUSet.swift

@@ -26,9 +26,6 @@ import CNIOLinux
         ///     - cpuIds: The `Set` of CPU ids. It must be non-empty and can not contain invalid ids.
         init(cpuIds: Set<Int>) {
             precondition(!cpuIds.isEmpty)
-            cpuIds.forEach{ v in
-                precondition(v >= 0 && v < System.coreCount)
-            }
             self.cpuIds = cpuIds
         }
 

Sources/NIO/Utilities.swift

@@ -79,7 +79,14 @@ public enum System {
             .filter { $0.Relationship == RelationProcessorCore }
             .map { $0.ProcessorMask.nonzeroBitCount }
             .reduce(0, +)
-
+#elseif os(Linux)
+        if let quota = Linux.coreCount(quota: Linux.cfsQuotaPath, period: Linux.cfsPeriodPath) {
+            return quota
+        } else if let cpusetCount = Linux.coreCount(cpuset: Linux.cpuSetPath) {
+            return cpusetCount
+        } else {
+            return sysconf(CInt(_SC_NPROCESSORS_ONLN))
+        }
 #else
         return sysconf(CInt(_SC_NPROCESSORS_ONLN))
 #endif

Tests/LinuxMain.swift

@@ -85,6 +85,7 @@ class LinuxMainRunnerImpl: LinuxMainRunner {
              testCase(IOErrorTest.allTests),
              testCase(IdleStateHandlerTest.allTests),
              testCase(IntegerTypesTest.allTests),
+             testCase(LinuxTest.allTests),
              testCase(MarkedCircularBufferTests.allTests),
              testCase(MessageToByteEncoderTest.allTests),
              testCase(MessageToByteHandlerTest.allTests),

Tests/NIOTests/LinuxTest+XCTest.swift

@@ -0,0 +1,35 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the SwiftNIO open source project
+//
+// Copyright (c) 2017-2018 Apple Inc. and the SwiftNIO project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of SwiftNIO project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+//
+// LinuxTest+XCTest.swift
+//
+import XCTest
+
+///
+/// NOTE: This file was generated by generate_linux_tests.rb
+///
+/// Do NOT edit this file directly as it will be regenerated automatically when needed.
+///
+
+extension LinuxTest {
+
+   @available(*, deprecated, message: "not actually deprecated. Just deprecated to allow deprecated tests (which test deprecated functionality) without warnings")
+   static var allTests : [(String, (LinuxTest) -> () throws -> Void)] {
+      return [
+                ("testCoreCountQuota", testCoreCountQuota),
+                ("testCoreCountCpuset", testCoreCountCpuset),
+           ]
+   }
+}
+

Tests/NIOTests/LinuxTest.swift

@@ -0,0 +1,60 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the SwiftNIO open source project
+//
+// Copyright (c) 2017-2020 Apple Inc. and the SwiftNIO project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of SwiftNIO project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+import XCTest
+@testable import NIO
+
+class LinuxTest: XCTestCase {
+    func testCoreCountQuota() {
+        #if os(Linux)
+        [
+            ("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)
+        ].forEach { quota, period, count in
+            withTemporaryFile(content: quota) { (_, quotaPath) -> Void in
+                withTemporaryFile(content: period) { (_, periodPath) -> Void in
+                    XCTAssertEqual(Linux.coreCount(quota: quotaPath, period: periodPath), count)
+                }
+            }
+        }
+        #endif
+    }
+
+    func testCoreCountCpuset() {
+        #if os(Linux)
+        [
+            ("0", 1),
+            ("0,3", 2),
+            ("0-3", 4),
+            ("0-3,7", 5),
+            ("0-3,7\n", 5),
+            ("0,2-4,6,7,9-11", 9),
+            ("", nil)
+        ].forEach { cpuset, count in
+            withTemporaryFile(content: cpuset) { (_, path) -> Void in
+                XCTAssertEqual(Linux.coreCount(cpuset: path), count)
+            }
+        }
+        #endif
+    }
+}