如何在Windows中使用Go将进程设置为CPU？

Question

I want set a process to a CPU using Go in win7, the below code:

package main

import (
    "fmt"
    "math"
    "runtime"
    "syscall"
    "unsafe"
)

func SetAffinity(pid int, mask *int64) {
    syscall.Syscall(syscall.SYS_SCHED_SETAFFINITY,
        uintptr(pid), 8, uintptr(unsafe.Pointer(mask)))
}

func GetAffinity(pid int, mask *int64) {
    syscall.Syscall(syscall.SYS_SCHED_GETAFFINITY,
        uintptr(pid), 8, uintptr(unsafe.Pointer(mask)))
}

var cpuNum = float64(runtime.NumCPU())

var setx = []struct {
    args     int
    expected int64
}{
    {0, int64(math.Pow(2, cpuNum)) - 2},
}

func main() {
    for _, ca := range setx {
        var cpuSet int64
        GetAffinity(ca.args, &cpuSet)
        cpuSet = cpuSet & 0XFFFFFFE
        SetAffinity(ca.args, &cpuSet)
        fmt.Println(cpuSet)
        GetAffinity(ca.args, &cpuSet)
        fmt.Println(cpuSet)
    }
}

When I use go run affinity.go, get the follow info:

# command-line-arguments
.\affinity.go:12: undefined: syscall.SYS_SCHED_SETAFFINITY
.\affinity.go:13: not enough arguments in call to syscall.Syscall
.\affinity.go:17: undefined: syscall.SYS_SCHED_GETAFFINITY
.\affinity.go:18: not enough arguments in call to syscall.Syscall

I find SYS_SCHED_SETAFFINITY that it only used in linux. So, I want to set a process to a cpu using Go in Windows(Win7), what can I do?

Answer 1

You'll have to invoke the WinAPI SetProcessAffinityMask.

Something like this should work:

func setProcessAffinityMask(h syscall.Handle, mask uintptr) (err error) {
    r1, _, e1 := syscall.Syscall(syscall.NewLazyDLL("kernel32.dll").NewProc("SetProcessAffinityMask").Addr(), 2, uintptr(h), mask, 0)
    if r1 == 0 {
        if e1 != 0 {
            err = error(e1)
        } else {
            err = syscall.EINVAL
        }
    }
    return
}

h being the process handle, and mask being the desired affinity mask, of course.

This is taken from Go benchmarks, under the BSD license.