affinity_linux.go

     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  }
   190
View as plain text