src

Go monorepo.
git clone git://code.dwrz.net/src
Log | Files | Refs

affinity_linux.go (5156B)

      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 bit mask of CPUs, to be used with [SchedGetaffinity], [SchedSetaffinity],
     17 // and [SetMemPolicy].
     18 //
     19 // Note this type can only represent CPU IDs 0 through 1023.
     20 // Use [CPUSetDynamic]/[NewCPUSet] instead to avoid this limit.
     21 type CPUSet [cpuSetSize]cpuMask
     22 
     23 // CPUSetDynamic represents a bit mask of CPUs, to be used with [SchedGetaffinityDynamic],
     24 // [SchedSetaffinityDynamic], and [SetMemPolicyDynamic]. Use [NewCPUSet] to allocate.
     25 type CPUSetDynamic []cpuMask
     26 
     27 func schedAffinity(trap uintptr, pid int, size uintptr, ptr unsafe.Pointer) error {
     28 	_, _, e := RawSyscall(trap, uintptr(pid), uintptr(size), uintptr(ptr))
     29 	if e != 0 {
     30 		return errnoErr(e)
     31 	}
     32 	return nil
     33 }
     34 
     35 // SchedGetaffinity gets the CPU affinity mask of the thread specified by pid.
     36 // If pid is 0 the calling thread is used.
     37 func SchedGetaffinity(pid int, set *CPUSet) error {
     38 	return schedAffinity(SYS_SCHED_GETAFFINITY, pid, unsafe.Sizeof(*set), unsafe.Pointer(set))
     39 }
     40 
     41 // SchedSetaffinity sets the CPU affinity mask of the thread specified by pid.
     42 // If pid is 0 the calling thread is used.
     43 func SchedSetaffinity(pid int, set *CPUSet) error {
     44 	return schedAffinity(SYS_SCHED_SETAFFINITY, pid, unsafe.Sizeof(*set), unsafe.Pointer(set))
     45 }
     46 
     47 // Zero clears the set s, so that it contains no CPUs.
     48 func (s *CPUSet) Zero() {
     49 	clear(s[:])
     50 }
     51 
     52 // Fill adds all possible CPU bits to the set s. On Linux, [SchedSetaffinity]
     53 // will silently ignore any invalid CPU bits in [CPUSet] so this is an
     54 // efficient way of resetting the CPU affinity of a process.
     55 func (s *CPUSet) Fill() {
     56 	cpuMaskFill(s[:])
     57 }
     58 
     59 func cpuBitsIndex(cpu int) int {
     60 	return cpu / _NCPUBITS
     61 }
     62 
     63 func cpuBitsMask(cpu int) cpuMask {
     64 	return cpuMask(1 << (uint(cpu) % _NCPUBITS))
     65 }
     66 
     67 func cpuMaskFill(s []cpuMask) {
     68 	for i := range s {
     69 		s[i] = ^cpuMask(0)
     70 	}
     71 }
     72 
     73 func cpuMaskSet(s []cpuMask, cpu int) {
     74 	i := cpuBitsIndex(cpu)
     75 	if i < len(s) {
     76 		s[i] |= cpuBitsMask(cpu)
     77 	}
     78 }
     79 
     80 func cpuMaskClear(s []cpuMask, cpu int) {
     81 	i := cpuBitsIndex(cpu)
     82 	if i < len(s) {
     83 		s[i] &^= cpuBitsMask(cpu)
     84 	}
     85 }
     86 
     87 func cpuMaskIsSet(s []cpuMask, cpu int) bool {
     88 	i := cpuBitsIndex(cpu)
     89 	if i < len(s) {
     90 		return s[i]&cpuBitsMask(cpu) != 0
     91 	}
     92 	return false
     93 }
     94 
     95 func cpuMaskCount(s []cpuMask) int {
     96 	c := 0
     97 	for _, b := range s {
     98 		c += bits.OnesCount64(uint64(b))
     99 	}
    100 	return c
    101 }
    102 
    103 // Set adds cpu to the set s. If cpu is out of bounds for s, no action is taken.
    104 func (s *CPUSet) Set(cpu int) {
    105 	cpuMaskSet(s[:], cpu)
    106 }
    107 
    108 // Clear removes cpu from the set s. If cpu is out of bounds for s, no action is taken.
    109 func (s *CPUSet) Clear(cpu int) {
    110 	cpuMaskClear(s[:], cpu)
    111 }
    112 
    113 // IsSet reports whether cpu is in the set s.
    114 func (s *CPUSet) IsSet(cpu int) bool {
    115 	return cpuMaskIsSet(s[:], cpu)
    116 }
    117 
    118 // Count returns the number of CPUs in the set s.
    119 func (s *CPUSet) Count() int {
    120 	return cpuMaskCount(s[:])
    121 }
    122 
    123 // NewCPUSet creates a CPU affinity mask capable of representing CPU IDs
    124 // up to maxCPU (exclusive).
    125 func NewCPUSet(maxCPU int) CPUSetDynamic {
    126 	numMasks := (maxCPU + _NCPUBITS - 1) / _NCPUBITS
    127 	if numMasks == 0 {
    128 		numMasks = 1
    129 	}
    130 	return make(CPUSetDynamic, numMasks)
    131 }
    132 
    133 // Zero clears the set s, so that it contains no CPUs.
    134 func (s CPUSetDynamic) Zero() {
    135 	clear(s)
    136 }
    137 
    138 // Fill adds all possible CPU bits to the set s. On Linux, [SchedSetaffinityDynamic]
    139 // will silently ignore any invalid CPU bits in [CPUSetDynamic] so this is an
    140 // efficient way of resetting the CPU affinity of a process.
    141 func (s CPUSetDynamic) Fill() {
    142 	cpuMaskFill(s)
    143 }
    144 
    145 // Set adds cpu to the set s. If cpu is out of bounds for s, no action is taken.
    146 func (s CPUSetDynamic) Set(cpu int) {
    147 	cpuMaskSet(s, cpu)
    148 }
    149 
    150 // Clear removes cpu from the set s. If cpu is out of bounds for s, no action is taken.
    151 func (s CPUSetDynamic) Clear(cpu int) {
    152 	cpuMaskClear(s, cpu)
    153 }
    154 
    155 // IsSet reports whether cpu is in the set s.
    156 func (s CPUSetDynamic) IsSet(cpu int) bool {
    157 	return cpuMaskIsSet(s, cpu)
    158 }
    159 
    160 // Count returns the number of CPUs in the set s.
    161 func (s CPUSetDynamic) Count() int {
    162 	return cpuMaskCount(s)
    163 }
    164 
    165 func (s CPUSetDynamic) size() uintptr {
    166 	return uintptr(len(s)) * unsafe.Sizeof(cpuMask(0))
    167 }
    168 
    169 func (s CPUSetDynamic) pointer() unsafe.Pointer {
    170 	if len(s) == 0 {
    171 		return nil
    172 	}
    173 	return unsafe.Pointer(&s[0])
    174 }
    175 
    176 // SchedGetaffinityDynamic gets the CPU affinity mask of the thread specified by pid.
    177 // If pid is 0 the calling thread is used.
    178 //
    179 // If the set is smaller than the size of the affinity mask used by the kernel,
    180 // [EINVAL] is returned.
    181 func SchedGetaffinityDynamic(pid int, set CPUSetDynamic) error {
    182 	return schedAffinity(SYS_SCHED_GETAFFINITY, pid, set.size(), set.pointer())
    183 }
    184 
    185 // SchedSetaffinityDynamic sets the CPU affinity mask of the thread specified by pid.
    186 // If pid is 0 the calling thread is used.
    187 func SchedSetaffinityDynamic(pid int, set CPUSetDynamic) error {
    188 	return schedAffinity(SYS_SCHED_SETAFFINITY, pid, set.size(), set.pointer())
    189 }