/src/swift-nio/Sources/NIOPosix/LinuxCPUSet.swift
Line | Count | Source |
1 | | //===----------------------------------------------------------------------===// |
2 | | // |
3 | | // This source file is part of the SwiftNIO open source project |
4 | | // |
5 | | // Copyright (c) 2017-2018 Apple Inc. and the SwiftNIO project authors |
6 | | // Licensed under Apache License v2.0 |
7 | | // |
8 | | // See LICENSE.txt for license information |
9 | | // See CONTRIBUTORS.txt for the list of SwiftNIO project authors |
10 | | // |
11 | | // SPDX-License-Identifier: Apache-2.0 |
12 | | // |
13 | | //===----------------------------------------------------------------------===// |
14 | | |
15 | | #if !os(WASI) |
16 | | |
17 | | #if os(Linux) || os(Android) |
18 | | import CNIOLinux |
19 | | |
20 | | /// A set that contains CPU ids to use. |
21 | | struct LinuxCPUSet { |
22 | | /// The ids of all the cpus. |
23 | | let cpuIds: Set<Int> |
24 | | |
25 | | /// Create a new instance |
26 | | /// |
27 | | /// - arguments: |
28 | | /// - cpuIds: The `Set` of CPU ids. It must be non-empty and can not contain invalid ids. |
29 | 0 | init(cpuIds: Set<Int>) { |
30 | 0 | precondition(!cpuIds.isEmpty) |
31 | 0 | self.cpuIds = cpuIds |
32 | 0 | } |
33 | | |
34 | | /// Create a new instance |
35 | | /// |
36 | | /// - arguments: |
37 | | /// - cpuId: The CPU id. |
38 | 0 | init(_ cpuId: Int) { |
39 | 0 | let ids: Set<Int> = [cpuId] |
40 | 0 | self.init(cpuIds: ids) |
41 | 0 | } |
42 | | } |
43 | | |
44 | | extension LinuxCPUSet: Equatable {} |
45 | | |
46 | | /// Linux specific extension to `NIOThread`. |
47 | | extension NIOThread { |
48 | | /// Specify the thread-affinity of the current thread. |
49 | | static var currentAffinity: LinuxCPUSet { |
50 | 0 | get { |
51 | 0 | var cpuset = cpu_set_t() |
52 | 0 |
|
53 | 0 | // Ensure the cpuset is empty (and so nothing is selected yet). |
54 | 0 | CNIOLinux_CPU_ZERO(&cpuset) |
55 | 0 |
|
56 | 0 | let res = CNIOLinux_pthread_getaffinity_np(pthread_self(), MemoryLayout.size(ofValue: cpuset), &cpuset) |
57 | 0 |
|
58 | 0 | precondition(res == 0, "pthread_getaffinity_np failed: \(res)") |
59 | 0 |
|
60 | 0 | let set = Set( |
61 | 0 | (CInt(0)..<CNIOLinux_CPU_SETSIZE()).lazy.filter { CNIOLinux_CPU_ISSET($0, &cpuset) != 0 }.map { |
62 | 0 | Int($0) |
63 | 0 | } |
64 | 0 | ) |
65 | 0 | return LinuxCPUSet(cpuIds: set) |
66 | 0 | } |
67 | 0 | set(cpuSet) { |
68 | 0 | var cpuset = cpu_set_t() |
69 | 0 |
|
70 | 0 | // Ensure the cpuset is empty (and so nothing is selected yet). |
71 | 0 | CNIOLinux_CPU_ZERO(&cpuset) |
72 | 0 |
|
73 | 0 | // Mark the CPU we want to run on. |
74 | 0 | for cpuID in cpuSet.cpuIds { |
75 | 0 | CNIOLinux_CPU_SET(CInt(cpuID), &cpuset) |
76 | 0 | } |
77 | 0 | let res = CNIOLinux_pthread_setaffinity_np(pthread_self(), MemoryLayout.size(ofValue: cpuset), &cpuset) |
78 | 0 | precondition(res == 0, "pthread_setaffinity_np failed: \(res)") |
79 | 0 | } |
80 | | } |
81 | | } |
82 | | |
83 | | extension MultiThreadedEventLoopGroup { |
84 | | |
85 | | /// Create a new `MultiThreadedEventLoopGroup` that create as many `NIOThread`s as `pinnedCPUIDs`. Each `NIOThread` will be pinned to the CPU with the id. |
86 | | /// |
87 | | /// Generally you should avoid using this API unless you have a concrete reason to need to pin threads to specific IDs. |
88 | | /// |
89 | | /// - arguments: |
90 | | /// - pinnedCPUIDs: The CPU ids to apply to the `NIOThread`s. |
91 | 0 | public convenience init(pinnedCPUIDs: [Int]) { |
92 | 0 | let initializers: [ThreadInitializer] = pinnedCPUIDs.map { id in |
93 | 0 | // This will also take care of validation of the provided id. |
94 | 0 | let set = LinuxCPUSet(id) |
95 | 0 | return { t in |
96 | 0 | precondition(t.isCurrentSlow) |
97 | 0 | NIOThread.currentAffinity = set |
98 | 0 | } |
99 | 0 | } |
100 | 0 | self.init(threadInitializers: initializers, metricsDelegate: nil) |
101 | 0 | } |
102 | | } |
103 | | #endif |
104 | | #endif // !os(WASI) |