src

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

affinity_linux.go (2189B)

      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 }