Go: x/sys/unix: add sched_getaffinity/sched_setaffinity

[wheelcomplex]

In Linux system resource manager programing(write in Go), we need sched_setaffinity to limit running CPU list of managed service process(no write in Go, C/C++/Java etc.) and use sched_getaffinity to verify it's corrected limited.

Export sched_getaffinity/sched_setaffinity func will be good for such purpose.

more in google group: https://groups.google.com/forum/#!topic/golang-nuts/99F1rtvcSSI

[adg]

cc @aclements @rsc @RLH @dr2chase

[aclements]

I think it would be a mistake to export these from the runtime for the reason that Ian mentioned on golang-nuts. You basically have to disable the Go scheduler in order to use these from Go, even for the purposes of writing a resource manager, and it would be strange to add something so incompatible with the runtime to the runtime API.

Ian mentioned adding these to x/sys. I don't know what the policy for x/sys is, but I assume this would be a fine collection of preexisting foot-guns to add these to. :)

Note that you can already get at these through syscall.RawSyscall using something like the following (untested):

func pinToCPU(cpu uint) error {
    const __NR_sched_setaffinity = 203
    var mask [1024 / 64]uint8
    runtime.LockOSThread()
    mask[cpu / 64] |= 1 << (cpu % 64)
    _, _, errno := syscall.RawSyscall(__NR_sched_setaffinity, 0, uintptr(len(mask) * 8), uintptr(unsafe.Pointer(&mask)))
    if errno != 0 {
        return errno
    }
    return nil
}

[i3d]

Should you call UnlockOSThread at the end of the func (or defer)?

[wheelcomplex]

@aclements thanks for the hint of RawSyscall usage, copied to google group.

To be clear, change the code to make sure it's using in Go program to setup affinity mask for a outside/standalone process(no write in Go).

// pinToCPU pin the pid to CPU index,
// usually, pid is a outside/standalone process(no write in Go) 
func pinToCPU(pid int, cpu uint) error {
    const __NR_sched_setaffinity = 203
    var mask [1024 / 64]uint8
    runtime.LockOSThread()
    defer runtime.UnlockOSThread()
    mask[cpu/64] |= 1 << (cpu % 64)
    _, _, errno := syscall.RawSyscall(__NR_sched_setaffinity, uintptr(pid), uintptr(len(mask)*8), uintptr(unsafe.Pointer(&mask)))
    if errno != 0 {
        return errno
    }
    return nil
}

[aclements]

Should you call UnlockOSThread at the end of the func (or defer)?

No. If it sets the affinity for the current OS thread and then unlocks the Goroutine from the OS thread, then the calling goroutine can be scheduled away from the affinitized OS thread and other goroutines can be scheduled on to it. This is one of the reasons I called this a foot gun.

[aclements]

@wheelcomplex: Ah, I hadn't appreciated that you wanted to change the affinity of an existing other PID. In that case, you don't need the Lock/UnlockOSThread at all. I had assumed you meant you were setting your own thread's affinity and then fork/exec'ing a new process that would inherit the CPU mask. Out of curiosity, how do you deal with setting the affinity of all of the other process's threads while the other process may be concurrently creating or destroying its threads?

[wheelcomplex]

@aclements We make sure outside/standalone process(write by another team and no in Go) can not creating or destroying its threads by administrator command(rule of work), no limit by code. In fact, we name outside/standalone process as "work load cell", there are many type of cells, they all code in single thread and launch or kill by resource manager(part of cluster manager) base on system load or business requirements changes.

[wheelcomplex]

Here is a small patch base on tip commit 71859efceb21a3f8098716ad8cf0964571be7bc5, only tested in ubuntu vivid amd64. It offer GetAffinity/SetAffinityAddAffinity for setup affinity of one pid(thread id in multi-thread process) and new field "Affinitys []int" in type SysProcAttr struct to setup affinity when using os/exec.Cmd. Hope this get Go better.

http://play.golang.org/p/EUUhXsOudU

package main

import "os/exec"

func main() {
    // exe a program to run on cpu 0, 1, 2
    cmd := &exec.Cmd{
        Path: "/usr/bin/stress",
        SysProcAttr: &syscall.SysProcAttr{
            Affinitys: []int{0, 1, 2},
        },
    }
    err := cmd.Start()
    if err != nil {
        panic(err.Error())
    }
    cmd.Wait()
}

The patch.

--- src/syscall/exec_linux.go.orig
+++ src/syscall/exec_linux.go
@@ -19,6 +19,7 @@
 }

 type SysProcAttr struct {
+   Affinitys   []int          // CPU affinitys.
    Chroot      string         // Chroot.
    Credential  *Credential    // Credential.
    Ptrace      bool           // Enable tracing.
@@ -44,6 +45,71 @@
 func runtime_BeforeFork()
 func runtime_AfterFork()

+const ptrSize = 4 << (^uintptr(0) >> 63) // unsafe.Sizeof(uintptr(0)) but an ideal const
+
+// GetAffinity return cpu list of pid
+func GetAffinity(pid uintptr) ([]int, error) {
+   var mask [1024 / 64]uintptr
+   var ret = make([]int, 0)
+   if pid <= 0 {
+       pid, _, _ = RawSyscall(SYS_GETPID, 0, 0, 0)
+   }
+   // size of uintptr is 8, in amd64
+   v1, _, err := RawSyscall(SYS_SCHED_GETAFFINITY, pid, uintptr(len(mask)*8), uintptr(unsafe.Pointer(&mask[0])))
+   if err != 0 {
+       return ret, err
+   }
+   nmask := mask[:v1/ptrSize]
+   idx := 0
+   for _, v := range nmask {
+       for i := 0; i < 64; i++ {
+           ct := int32(v & 1)
+           v >>= 1
+           if ct > 0 {
+               ret = append(ret, idx)
+           }
+           idx++
+       }
+   }
+   return ret, nil
+}
+
+// SetAffinity attend the cpu list to pid,
+// note: SetAffinity apply to thread ID only,
+// to fully control one process, call SetAffinity for all thread of the process.
+// use os.GetThreadIDs() to get all thread of the process
+func SetAffinity(pid uintptr, cpus []int) error {
+   var mask [1024 / 64]uintptr
+   if pid <= 0 {
+       pid, _, _ = RawSyscall(SYS_GETPID, 0, 0, 0)
+   }
+   for _, cpuIdx := range cpus {
+       cpuIndex := uint(cpuIdx)
+       mask[cpuIndex/64] |= 1 << (cpuIndex % 64)
+   }
+   _, _, err := RawSyscall(SYS_SCHED_SETAFFINITY, pid, uintptr(len(mask)*8), uintptr(unsafe.Pointer(&mask[0])))
+   if err != 0 {
+       return err
+   }
+   return nil
+}
+
+// AddAffinity attend one cpu to list of pid
+// note: AddAffinity apply to thread ID only,
+// to fully control one process, call SetAffinity for all thread of the process
+// use os.GetThreadIDs() to get all thread of the process
+func AddAffinity(pid uintptr, cpuIdx int) error {
+   if pid <= 0 {
+       pid, _, _ = RawSyscall(SYS_GETPID, 0, 0, 0)
+   }
+   cpus, e1 := GetAffinity(pid)
+   if e1 != nil {
+       return e1
+   }
+   cpus = append(cpus, cpuIdx)
+   return SetAffinity(pid, cpus)
+}
+
 // Fork, dup fd onto 0..len(fd), and exec(argv0, argvv, envv) in child.
 // If a dup or exec fails, write the errno error to pipe.
 // (Pipe is close-on-exec so if exec succeeds, it will be closed.)
@@ -63,6 +129,7 @@
        nextfd int
        i      int
        p      [2]int
+       mask   [1024 / 64]uintptr
    )

    // Record parent PID so child can test if it has died.
@@ -117,6 +184,19 @@
    }

    // Fork succeeded, now in child.
+
+   // setup cpu affinity when required
+   if sys.Affinitys != nil {
+       upid, _, _ := RawSyscall(SYS_GETPID, 0, 0, 0)
+       for _, cpuIdx := range sys.Affinitys {
+           cpuIndex := uint(cpuIdx)
+           mask[cpuIndex/64] |= 1 << (cpuIndex % 64)
+       }
+       _, _, err1 = RawSyscall(SYS_SCHED_SETAFFINITY, upid, uintptr(len(mask)*8), uintptr(unsafe.Pointer(&mask[0])))
+       if err1 != 0 {
+           goto childerror
+       }
+   }

    // Wait for User ID/Group ID mappings to be written.
    if sys.UidMappings != nil || sys.GidMappings != nil {

[ianlancetaylor]

The general rule for SysProcAttr is that it is only used for operations that must occur between fork and exec. That is not the case for sched_setaffinity. In particular, you can do the operation by simply changing your command to invoke the taskset program. So that change should not be made.

As noted previously, the SetAffinity, etc., functions should not be added to the syscall package, because that package is frozen. They should be added to golang.org/x/sys/unix.

[caoruidong]

sched_set/getaffinity are still not added to golang.org/x/sys/unix now?

[ianlancetaylor]

@caoruidong Correct. This issue is still open. Want to send a patch?

[caoruidong]

@ianlancetaylor OK. I will try

[gopherbot]

Change https://golang.org/cl/85915 mentions this issue: unix: add SchedGetaffinity and SchedSetaffinity on Linux

[dt-rush]

Edit: this comment only valid for go < 1.10

This is just a note for future readers of this thread. If you use the code provided by aclements in their Jun 17, 2015 comment, you will set the affinity for the thread which the goroutine has been locked to. But there's a problem - while the runtime has locked other goroutines from spawning on that thread - the scheduler might run on that thread and might decide to spawn other threads, which will inherit the affinity (the link talks about inheriting priority, but affinity masks are also inherited by threads). So you are then obliged to add lines at the top of every goroutine you spawn in your program to try to set the affinity mask to match all CPUs (AKA "unset"). You won't be able to add this for any libraries you import, so it's a bigger footgun than it seems, and basically not usable. We would have to add a call parallel to runtime.LockOSThread() which would somehow tell the runtime to not run the scheduler on the locked thread.

[aclements]

You're correct that this was an issue before Go 1.10, but as of Go 1.10 the
runtime specifically does not spawn new threads from locked threads
precisely because of issues like this, so this shouldn't be a problem.

[dt-rush]

I edited my comment to make sure nobody stops reading there. How awesome!!!!!!! I'm gonna start using this trick then hahaha