golang：设置 Win10 进程亲和性

1. 背景

文章中说明了如何设置进程的优先级。这篇文章说一下如何设置进程的亲和性。

亲和性是指进程和 CPU 核的亲和性。设置进程的亲和性后，程序只运行在设置的那几个核中，不会被调度到以外的核。

2. SetProcessAffinityMask() 说明

Win10 环境下，设置进程亲和性的接口为SetProcessAffinityMask()，函数原型如下：

BOOL SetProcessAffinityMask(
  [in] HANDLE    hProcess,
  [in] DWORD_PTR dwProcessAffinityMask
);

但是，这个函数在 go 中没有被实现，需要自己封装。

3. 接口封装及测试

比如说，想让进程可以运行在 0、3、6 核上。

package main

import "time"
import "fmt"
import "syscall"
import "golang.org/x/sys/windows"

//封装接口
var modkernel32 = windows.NewLazySystemDLL("kernel32.dll")
var procSetProcessAffinityMask = modkernel32.NewProc("SetProcessAffinityMask")

func SetProcessAffinityMask(s windows.Handle, processAffinityMask uint64) (err error) {
    r1, _, e1 := syscall.Syscall(procSetProcessAffinityMask.Addr(), 2, uintptr(s), uintptr(processAffinityMask), 0)
    if r1 == 0 {
        if e1 != 0 {
            err = e1
        } else {
            err = syscall.EINVAL
        }
    }
    return
}

func main() {
    cur_task_handle := windows.CurrentProcess()
    err := windows.SetPriorityClass(cur_task_handle, windows.HIGH_PRIORITY_CLASS)
    if err != nil {
        fmt.Println(err)
        return
    }

    err = SetProcessAffinityMask(cur_task_handle, ((1 << 0) | (1 << 3) | (1 << 6)))
    if err != nil {
        fmt.Println(err)
    }

    for  {
        time.Sleep(10 * time.Second)
    }
}

重点说明一下参数 processAffinityMask，这个参数是位有效，某一位置 1 表示进程可以运行在该 CPU 核，为 0 表示进程不能运行在该核。就像上述代码中，该参数的值为 ((1 << 0) | (1 << 3) | (1 << 6))，就表示程序可以运行在 0、3、6 核。

4. 总结

如果想要设置线程的亲和性，也可以按照上述方式进行封装。

import "syscall"
import "golang.org/x/sys/windows"

var modkernel32 = windows.NewLazySystemDLL("kernel32.dll")
var procSetThreadAffinityMask = modkernel32.NewProc("SetThreadAffinityMask")
var procSetProcessAffinityMask = modkernel32.NewProc("SetProcessAffinityMask")

func SetThreadAffinityMask(s windows.Handle, threadAffinityMask uint64) (err error) {
    r1, _, e1 := syscall.Syscall(procSetThreadAffinityMask.Addr(), 2, uintptr(s), uintptr(threadAffinityMask), 0)
    if r1 == 0 {
        if e1 != 0 {
            err = e1
        } else {
            err = syscall.EINVAL
        }
    }
    return
}