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
 7package unix
 8
 9import (
10	"math/bits"
11	"unsafe"
12)
13
14const cpuSetSize = _CPU_SETSIZE / _NCPUBITS
15
16// CPUSet represents a CPU affinity mask.
17type CPUSet [cpuSetSize]cpuMask
18
19func 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.
29func 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.
35func 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.
40func (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.
47func (s *CPUSet) Fill() {
48	for i := range s {
49		s[i] = ^cpuMask(0)
50	}
51}
52
53func cpuBitsIndex(cpu int) int {
54	return cpu / _NCPUBITS
55}
56
57func cpuBitsMask(cpu int) cpuMask {
58	return cpuMask(1 << (uint(cpu) % _NCPUBITS))
59}
60
61// Set adds cpu to the set s.
62func (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.
70func (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.
78func (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.
87func (s *CPUSet) Count() int {
88	c := 0
89	for _, b := range s {
90		c += bits.OnesCount64(uint64(b))
91	}
92	return c
93}