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

95