1 // Copyright 2018 The Go Authors. All rights reserved.
2 // Use of this source code is governed by a BSD-style
3 // license that can be found in the LICENSE file.
4 5 // CPU affinity functions
6 7 package unix
8 9 import (
10 "math/bits"
11 "unsafe"
12 )
13 14 const cpuSetSize = _CPU_SETSIZE / _NCPUBITS
15 16 // CPUSet represents a CPU affinity mask.
17 type CPUSet [cpuSetSize]cpuMask
18 19 func schedAffinity(trap uintptr, pid int, set *CPUSet) error {
20 _, _, e := RawSyscall(trap, uintptr(pid), uintptr(unsafe.Sizeof(*set)), uintptr(unsafe.Pointer(set)))
21 if e != 0 {
22 return errnoErr(e)
23 }
24 return nil
25 }
26 27 // SchedGetaffinity gets the CPU affinity mask of the thread specified by pid.
28 // If pid is 0 the calling thread is used.
29 func SchedGetaffinity(pid int, set *CPUSet) error {
30 return schedAffinity(SYS_SCHED_GETAFFINITY, pid, set)
31 }
32 33 // SchedSetaffinity sets the CPU affinity mask of the thread specified by pid.
34 // If pid is 0 the calling thread is used.
35 func SchedSetaffinity(pid int, set *CPUSet) error {
36 return schedAffinity(SYS_SCHED_SETAFFINITY, pid, set)
37 }
38 39 // Zero clears the set s, so that it contains no CPUs.
40 func (s *CPUSet) Zero() {
41 clear(s[:])
42 }
43 44 // Fill adds all possible CPU bits to the set s. On Linux, [SchedSetaffinity]
45 // will silently ignore any invalid CPU bits in [CPUSet] so this is an
46 // efficient way of resetting the CPU affinity of a process.
47 func (s *CPUSet) Fill() {
48 for i := range s {
49 s[i] = ^cpuMask(0)
50 }
51 }
52 53 func cpuBitsIndex(cpu int) int {
54 return cpu / _NCPUBITS
55 }
56 57 func cpuBitsMask(cpu int) cpuMask {
58 return cpuMask(1 << (uint(cpu) % _NCPUBITS))
59 }
60 61 // Set adds cpu to the set s.
62 func (s *CPUSet) Set(cpu int) {
63 i := cpuBitsIndex(cpu)
64 if i < len(s) {
65 s[i] |= cpuBitsMask(cpu)
66 }
67 }
68 69 // Clear removes cpu from the set s.
70 func (s *CPUSet) Clear(cpu int) {
71 i := cpuBitsIndex(cpu)
72 if i < len(s) {
73 s[i] &^= cpuBitsMask(cpu)
74 }
75 }
76 77 // IsSet reports whether cpu is in the set s.
78 func (s *CPUSet) IsSet(cpu int) bool {
79 i := cpuBitsIndex(cpu)
80 if i < len(s) {
81 return s[i]&cpuBitsMask(cpu) != 0
82 }
83 return false
84 }
85 86 // Count returns the number of CPUs in the set s.
87 func (s *CPUSet) Count() int {
88 c := 0
89 for _, b := range s {
90 c += bits.OnesCount64(uint64(b))
91 }
92 return c
93 }
94