affinity_linux.go 1.9 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586
  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. // CPU affinity functions
  5. package unix
  6. import (
  7. "math/bits"
  8. "unsafe"
  9. )
  10. const cpuSetSize = _CPU_SETSIZE / _NCPUBITS
  11. // CPUSet represents a CPU affinity mask.
  12. type CPUSet [cpuSetSize]cpuMask
  13. func schedAffinity(trap uintptr, pid int, set *CPUSet) error {
  14. _, _, e := RawSyscall(trap, uintptr(pid), uintptr(unsafe.Sizeof(*set)), uintptr(unsafe.Pointer(set)))
  15. if e != 0 {
  16. return errnoErr(e)
  17. }
  18. return nil
  19. }
  20. // SchedGetaffinity gets the CPU affinity mask of the thread specified by pid.
  21. // If pid is 0 the calling thread is used.
  22. func SchedGetaffinity(pid int, set *CPUSet) error {
  23. return schedAffinity(SYS_SCHED_GETAFFINITY, pid, set)
  24. }
  25. // SchedSetaffinity sets the CPU affinity mask of the thread specified by pid.
  26. // If pid is 0 the calling thread is used.
  27. func SchedSetaffinity(pid int, set *CPUSet) error {
  28. return schedAffinity(SYS_SCHED_SETAFFINITY, pid, set)
  29. }
  30. // Zero clears the set s, so that it contains no CPUs.
  31. func (s *CPUSet) Zero() {
  32. for i := range s {
  33. s[i] = 0
  34. }
  35. }
  36. func cpuBitsIndex(cpu int) int {
  37. return cpu / _NCPUBITS
  38. }
  39. func cpuBitsMask(cpu int) cpuMask {
  40. return cpuMask(1 << (uint(cpu) % _NCPUBITS))
  41. }
  42. // Set adds cpu to the set s.
  43. func (s *CPUSet) Set(cpu int) {
  44. i := cpuBitsIndex(cpu)
  45. if i < len(s) {
  46. s[i] |= cpuBitsMask(cpu)
  47. }
  48. }
  49. // Clear removes cpu from the set s.
  50. func (s *CPUSet) Clear(cpu int) {
  51. i := cpuBitsIndex(cpu)
  52. if i < len(s) {
  53. s[i] &^= cpuBitsMask(cpu)
  54. }
  55. }
  56. // IsSet reports whether cpu is in the set s.
  57. func (s *CPUSet) IsSet(cpu int) bool {
  58. i := cpuBitsIndex(cpu)
  59. if i < len(s) {
  60. return s[i]&cpuBitsMask(cpu) != 0
  61. }
  62. return false
  63. }
  64. // Count returns the number of CPUs in the set s.
  65. func (s *CPUSet) Count() int {
  66. c := 0
  67. for _, b := range s {
  68. c += bits.OnesCount64(uint64(b))
  69. }
  70. return c
  71. }