diff --git a/unix/affinity_linux.go b/unix/affinity_linux.go new file mode 100644 index 00000000..c7ed1bdd --- /dev/null +++ b/unix/affinity_linux.go @@ -0,0 +1,86 @@ +// Copyright 2018 The Go Authors. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +// CPU affinity functions + +package unix + +import ( + "math/bits" + "unsafe" +) + +const cpuSetSize = _CPU_SETSIZE / _NCPUBITS + +// CPUSet represents a CPU affinity mask. +type CPUSet [cpuSetSize]cpuMask + +func schedAffinity(trap uintptr, pid int, set *CPUSet) error { + _, _, e := RawSyscall(trap, uintptr(pid), uintptr(unsafe.Sizeof(set)), uintptr(unsafe.Pointer(set))) + if e != 0 { + return errnoErr(e) + } + return nil +} + +// SchedGetaffinity gets the CPU affinity mask of the thread specified by pid. +// If pid is 0 the calling thread is used. +func SchedGetaffinity(pid int, set *CPUSet) error { + return schedAffinity(SYS_SCHED_GETAFFINITY, pid, set) +} + +// SchedSetaffinity sets the CPU affinity mask of the thread specified by pid. +// If pid is 0 the calling thread is used. +func SchedSetaffinity(pid int, set *CPUSet) error { + return schedAffinity(SYS_SCHED_SETAFFINITY, pid, set) +} + +// Zero clears the set s, so that it contains no CPUs. +func (s *CPUSet) Zero() { + for i := range s { + s[i] = 0 + } +} + +func cpuBitsIndex(cpu int) int { + return cpu / _NCPUBITS +} + +func cpuBitsMask(cpu int) cpuMask { + return cpuMask(1 << (uint(cpu) % _NCPUBITS)) +} + +// Set adds cpu to the set s. +func (s *CPUSet) Set(cpu int) { + i := cpuBitsIndex(cpu) + if i < len(s) { + s[i] |= cpuBitsMask(cpu) + } +} + +// Clear removes cpu from the set s. +func (s *CPUSet) Clear(cpu int) { + i := cpuBitsIndex(cpu) + if i < len(s) { + s[i] &^= cpuBitsMask(cpu) + } +} + +// IsSet reports whether cpu is in the set s. +func (s *CPUSet) IsSet(cpu int) bool { + i := cpuBitsIndex(cpu) + if i < len(s) { + return s[i]&cpuBitsMask(cpu) != 0 + } + return false +} + +// Count returns the number of CPUs in the set s. +func (s *CPUSet) Count() int { + c := 0 + for _, b := range s { + c += bits.OnesCount64(uint64(b)) + } + return c +} diff --git a/unix/linux/types.go b/unix/linux/types.go index b4350f8b..e233ebc3 100644 --- a/unix/linux/types.go +++ b/unix/linux/types.go @@ -24,6 +24,7 @@ package unix #include #include #include +#include #include #include #include @@ -594,3 +595,12 @@ const ( CTRL_ATTR_MCAST_GRP_NAME = C.CTRL_ATTR_MCAST_GRP_NAME CTRL_ATTR_MCAST_GRP_ID = C.CTRL_ATTR_MCAST_GRP_ID ) + +// CPU affinity + +type cpuMask C.__cpu_mask + +const ( + _CPU_SETSIZE = C.__CPU_SETSIZE + _NCPUBITS = C.__NCPUBITS +) diff --git a/unix/syscall_linux.go b/unix/syscall_linux.go index f77f44a5..44628566 100644 --- a/unix/syscall_linux.go +++ b/unix/syscall_linux.go @@ -1455,11 +1455,9 @@ func Vmsplice(fd int, iovs []Iovec, flags int) (int, error) { // RtSigtimedwait // SchedGetPriorityMax // SchedGetPriorityMin -// SchedGetaffinity // SchedGetparam // SchedGetscheduler // SchedRrGetInterval -// SchedSetaffinity // SchedSetparam // SchedYield // Security diff --git a/unix/syscall_linux_test.go b/unix/syscall_linux_test.go index ea9562fe..31ae24c3 100644 --- a/unix/syscall_linux_test.go +++ b/unix/syscall_linux_test.go @@ -9,6 +9,7 @@ package unix_test import ( "io/ioutil" "os" + "runtime" "testing" "time" @@ -255,6 +256,58 @@ func TestFstatat(t *testing.T) { } } +func TestSchedSetaffinity(t *testing.T) { + runtime.LockOSThread() + defer runtime.UnlockOSThread() + + var oldMask unix.CPUSet + err := unix.SchedGetaffinity(0, &oldMask) + if err != nil { + t.Fatalf("SchedGetaffinity: %v", err) + } + + var newMask unix.CPUSet + newMask.Zero() + if newMask.Count() != 0 { + t.Errorf("CpuZero: didn't zero CPU set: %v", newMask) + } + cpu := 1 + newMask.Set(cpu) + if newMask.Count() != 1 || !newMask.IsSet(cpu) { + t.Errorf("CpuSet: didn't set CPU %d in set: %v", cpu, newMask) + } + cpu = 5 + newMask.Set(cpu) + if newMask.Count() != 2 || !newMask.IsSet(cpu) { + t.Errorf("CpuSet: didn't set CPU %d in set: %v", cpu, newMask) + } + newMask.Clear(cpu) + if newMask.Count() != 1 || newMask.IsSet(cpu) { + t.Errorf("CpuClr: didn't clear CPU %d in set: %v", cpu, newMask) + } + + err = unix.SchedSetaffinity(0, &newMask) + if err != nil { + t.Fatalf("SchedSetaffinity: %v", err) + } + + var gotMask unix.CPUSet + err = unix.SchedGetaffinity(0, &gotMask) + if err != nil { + t.Fatalf("SchedGetaffinity: %v", err) + } + + if gotMask != newMask { + t.Errorf("SchedSetaffinity: returned affinity mask does not match set affinity mask") + } + + // Restore old mask so it doesn't affect successive tests + err = unix.SchedSetaffinity(0, &oldMask) + if err != nil { + t.Fatalf("SchedSetaffinity: %v", err) + } +} + // utilities taken from os/os_test.go func touch(t *testing.T, name string) { diff --git a/unix/ztypes_linux_386.go b/unix/ztypes_linux_386.go index 295a9bd3..8f755b7f 100644 --- a/unix/ztypes_linux_386.go +++ b/unix/ztypes_linux_386.go @@ -792,3 +792,10 @@ const ( CTRL_ATTR_MCAST_GRP_NAME = 0x1 CTRL_ATTR_MCAST_GRP_ID = 0x2 ) + +type cpuMask uint32 + +const ( + _CPU_SETSIZE = 0x400 + _NCPUBITS = 0x20 +) diff --git a/unix/ztypes_linux_amd64.go b/unix/ztypes_linux_amd64.go index cdd4a1dc..92cfeb30 100644 --- a/unix/ztypes_linux_amd64.go +++ b/unix/ztypes_linux_amd64.go @@ -810,3 +810,10 @@ const ( CTRL_ATTR_MCAST_GRP_NAME = 0x1 CTRL_ATTR_MCAST_GRP_ID = 0x2 ) + +type cpuMask uint64 + +const ( + _CPU_SETSIZE = 0x400 + _NCPUBITS = 0x40 +) diff --git a/unix/ztypes_linux_arm.go b/unix/ztypes_linux_arm.go index 086620ac..9467efe3 100644 --- a/unix/ztypes_linux_arm.go +++ b/unix/ztypes_linux_arm.go @@ -781,3 +781,10 @@ const ( CTRL_ATTR_MCAST_GRP_NAME = 0x1 CTRL_ATTR_MCAST_GRP_ID = 0x2 ) + +type cpuMask uint32 + +const ( + _CPU_SETSIZE = 0x400 + _NCPUBITS = 0x20 +) diff --git a/unix/ztypes_linux_arm64.go b/unix/ztypes_linux_arm64.go index c53db98d..6e6def26 100644 --- a/unix/ztypes_linux_arm64.go +++ b/unix/ztypes_linux_arm64.go @@ -789,3 +789,10 @@ const ( CTRL_ATTR_MCAST_GRP_NAME = 0x1 CTRL_ATTR_MCAST_GRP_ID = 0x2 ) + +type cpuMask uint64 + +const ( + _CPU_SETSIZE = 0x400 + _NCPUBITS = 0x40 +) diff --git a/unix/ztypes_linux_mips.go b/unix/ztypes_linux_mips.go index 69e529b2..6d9094ff 100644 --- a/unix/ztypes_linux_mips.go +++ b/unix/ztypes_linux_mips.go @@ -786,3 +786,10 @@ const ( CTRL_ATTR_MCAST_GRP_NAME = 0x1 CTRL_ATTR_MCAST_GRP_ID = 0x2 ) + +type cpuMask uint32 + +const ( + _CPU_SETSIZE = 0x400 + _NCPUBITS = 0x20 +) diff --git a/unix/ztypes_linux_mips64.go b/unix/ztypes_linux_mips64.go index 4460279a..b169ff58 100644 --- a/unix/ztypes_linux_mips64.go +++ b/unix/ztypes_linux_mips64.go @@ -791,3 +791,10 @@ const ( CTRL_ATTR_MCAST_GRP_NAME = 0x1 CTRL_ATTR_MCAST_GRP_ID = 0x2 ) + +type cpuMask uint64 + +const ( + _CPU_SETSIZE = 0x400 + _NCPUBITS = 0x40 +) diff --git a/unix/ztypes_linux_mips64le.go b/unix/ztypes_linux_mips64le.go index 98923811..ec0aeb5d 100644 --- a/unix/ztypes_linux_mips64le.go +++ b/unix/ztypes_linux_mips64le.go @@ -791,3 +791,10 @@ const ( CTRL_ATTR_MCAST_GRP_NAME = 0x1 CTRL_ATTR_MCAST_GRP_ID = 0x2 ) + +type cpuMask uint64 + +const ( + _CPU_SETSIZE = 0x400 + _NCPUBITS = 0x40 +) diff --git a/unix/ztypes_linux_mipsle.go b/unix/ztypes_linux_mipsle.go index c2f4c036..0a656363 100644 --- a/unix/ztypes_linux_mipsle.go +++ b/unix/ztypes_linux_mipsle.go @@ -786,3 +786,10 @@ const ( CTRL_ATTR_MCAST_GRP_NAME = 0x1 CTRL_ATTR_MCAST_GRP_ID = 0x2 ) + +type cpuMask uint32 + +const ( + _CPU_SETSIZE = 0x400 + _NCPUBITS = 0x20 +) diff --git a/unix/ztypes_linux_ppc64.go b/unix/ztypes_linux_ppc64.go index ffe78833..869847db 100644 --- a/unix/ztypes_linux_ppc64.go +++ b/unix/ztypes_linux_ppc64.go @@ -799,3 +799,10 @@ const ( CTRL_ATTR_MCAST_GRP_NAME = 0x1 CTRL_ATTR_MCAST_GRP_ID = 0x2 ) + +type cpuMask uint64 + +const ( + _CPU_SETSIZE = 0x400 + _NCPUBITS = 0x40 +) diff --git a/unix/ztypes_linux_ppc64le.go b/unix/ztypes_linux_ppc64le.go index 7e3d946a..ddf720e3 100644 --- a/unix/ztypes_linux_ppc64le.go +++ b/unix/ztypes_linux_ppc64le.go @@ -799,3 +799,10 @@ const ( CTRL_ATTR_MCAST_GRP_NAME = 0x1 CTRL_ATTR_MCAST_GRP_ID = 0x2 ) + +type cpuMask uint64 + +const ( + _CPU_SETSIZE = 0x400 + _NCPUBITS = 0x40 +) diff --git a/unix/ztypes_linux_s390x.go b/unix/ztypes_linux_s390x.go index 730fa8a6..9ddb4470 100644 --- a/unix/ztypes_linux_s390x.go +++ b/unix/ztypes_linux_s390x.go @@ -816,3 +816,10 @@ const ( CTRL_ATTR_MCAST_GRP_NAME = 0x1 CTRL_ATTR_MCAST_GRP_ID = 0x2 ) + +type cpuMask uint64 + +const ( + _CPU_SETSIZE = 0x400 + _NCPUBITS = 0x40 +)