runtime: account for cpu affinity in windows NumCPU

Fixes #11671

Change-Id: Ide1f8d92637dad2a2faed391329f9b6001789b76
Reviewed-on: https://go-review.googlesource.com/14742
Reviewed-by: Russ Cox <rsc@golang.org>
Run-TryBot: Alex Brainman <alex.brainman@gmail.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>

diff --git a/src/runtime/export_windows_test.go b/src/runtime/export_windows_test.go
index 6322ee2804af272447bdb2539a9fdae261f6c28f..703b422ac2315abc50a33c46dc1059a892e90cc9 100644 (file)
--- a/src/runtime/export_windows_test.go
+++ b/src/runtime/export_windows_test.go
@@ -6,7 +6,15 @@
 
 package runtime
 
+import "unsafe"
+
 var (
        TestingWER              = &testingWER
        TimeBeginPeriodRetValue = &timeBeginPeriodRetValue
 )
+
+func NumberOfProcessors() int32 {
+       var info systeminfo
+       stdcall1(_GetSystemInfo, uintptr(unsafe.Pointer(&info)))
+       return int32(info.dwnumberofprocessors)
+}

diff --git a/src/runtime/os1_windows.go b/src/runtime/os1_windows.go
index b6da4df7c72e2b91e0f077bf62b73bee01e4b331..99c6df4008b028b22a21ee8eb256559248bca2e2 100644 (file)
--- a/src/runtime/os1_windows.go
+++ b/src/runtime/os1_windows.go
@@ -22,6 +22,7 @@ import (
 //go:cgo_import_dynamic runtime._FreeEnvironmentStringsW FreeEnvironmentStringsW%1 "kernel32.dll"
 //go:cgo_import_dynamic runtime._GetEnvironmentStringsW GetEnvironmentStringsW%0 "kernel32.dll"
 //go:cgo_import_dynamic runtime._GetProcAddress GetProcAddress%2 "kernel32.dll"
+//go:cgo_import_dynamic runtime._GetProcessAffinityMask GetProcessAffinityMask%3 "kernel32.dll"
 //go:cgo_import_dynamic runtime._GetQueuedCompletionStatus GetQueuedCompletionStatus%5 "kernel32.dll"
 //go:cgo_import_dynamic runtime._GetStdHandle GetStdHandle%1 "kernel32.dll"
 //go:cgo_import_dynamic runtime._GetSystemInfo GetSystemInfo%1 "kernel32.dll"
@@ -63,6 +64,7 @@ var (
        _FreeEnvironmentStringsW,
        _GetEnvironmentStringsW,
        _GetProcAddress,
+       _GetProcessAffinityMask,
        _GetQueuedCompletionStatus,
        _GetStdHandle,
        _GetSystemInfo,
@@ -126,6 +128,21 @@ func getGetProcAddress() uintptr {
 }
 
 func getproccount() int32 {
+       var mask, sysmask uintptr
+       ret := stdcall3(_GetProcessAffinityMask, currentProcess, uintptr(unsafe.Pointer(&mask)), uintptr(unsafe.Pointer(&sysmask)))
+       if ret != 0 {
+               n := 0
+               maskbits := int(unsafe.Sizeof(mask) * 8)
+               for i := 0; i < maskbits; i++ {
+                       if mask&(1<<uint(i)) != 0 {
+                               n++
+                       }
+               }
+               if n != 0 {
+                       return int32(n)
+               }
+       }
+       // use GetSystemInfo if GetProcessAffinityMask fails
        var info systeminfo
        stdcall1(_GetSystemInfo, uintptr(unsafe.Pointer(&info)))
        return int32(info.dwnumberofprocessors)

diff --git a/src/runtime/syscall_windows_test.go b/src/runtime/syscall_windows_test.go
index 4bedd4add4686c9ba7415053263c3d76f4ff61b6..8b9945b22ddfdf64c7e21baae7ad3c28cda1aab6 100644 (file)
--- a/src/runtime/syscall_windows_test.go
+++ b/src/runtime/syscall_windows_test.go
@@ -5,6 +5,7 @@
 package runtime_test
 
 import (
+       "bytes"
        "fmt"
        "io/ioutil"
        "os"
@@ -647,3 +648,155 @@ func TestTimeBeginPeriod(t *testing.T) {
                t.Fatalf("timeBeginPeriod failed: it returned %d", *runtime.TimeBeginPeriodRetValue)
        }
 }
+
+// removeOneCPU removes one (any) cpu from affinity mask.
+// It returns new affinity mask.
+func removeOneCPU(mask uintptr) (uintptr, error) {
+       if mask == 0 {
+               return 0, fmt.Errorf("cpu affinity mask is empty")
+       }
+       maskbits := int(unsafe.Sizeof(mask) * 8)
+       for i := 0; i < maskbits; i++ {
+               newmask := mask & ^(1 << uint(i))
+               if newmask != mask {
+                       return newmask, nil
+               }
+
+       }
+       panic("not reached")
+}
+
+func resumeChildThread(kernel32 *syscall.DLL, childpid int) error {
+       _OpenThread := kernel32.MustFindProc("OpenThread")
+       _ResumeThread := kernel32.MustFindProc("ResumeThread")
+       _Thread32First := kernel32.MustFindProc("Thread32First")
+       _Thread32Next := kernel32.MustFindProc("Thread32Next")
+
+       snapshot, err := syscall.CreateToolhelp32Snapshot(syscall.TH32CS_SNAPTHREAD, 0)
+       if err != nil {
+               return err
+       }
+       defer syscall.CloseHandle(snapshot)
+
+       const _THREAD_SUSPEND_RESUME = 0x0002
+
+       type ThreadEntry32 struct {
+               Size           uint32
+               tUsage         uint32
+               ThreadID       uint32
+               OwnerProcessID uint32
+               BasePri        int32
+               DeltaPri       int32
+               Flags          uint32
+       }
+
+       var te ThreadEntry32
+       te.Size = uint32(unsafe.Sizeof(te))
+       ret, _, err := _Thread32First.Call(uintptr(snapshot), uintptr(unsafe.Pointer(&te)))
+       if ret == 0 {
+               return err
+       }
+       for te.OwnerProcessID != uint32(childpid) {
+               ret, _, err = _Thread32Next.Call(uintptr(snapshot), uintptr(unsafe.Pointer(&te)))
+               if ret == 0 {
+                       return err
+               }
+       }
+       h, _, err := _OpenThread.Call(_THREAD_SUSPEND_RESUME, 1, uintptr(te.ThreadID))
+       if h == 0 {
+               return err
+       }
+       defer syscall.Close(syscall.Handle(h))
+
+       ret, _, err = _ResumeThread.Call(h)
+       if ret == 0xffffffff {
+               return err
+       }
+       return nil
+}
+
+func TestNumCPU(t *testing.T) {
+       if os.Getenv("GO_WANT_HELPER_PROCESS") == "1" {
+               // in child process
+               fmt.Fprintf(os.Stderr, "%d", runtime.NumCPU())
+               os.Exit(0)
+       }
+
+       switch n := runtime.NumberOfProcessors(); {
+       case n < 1:
+               t.Fatalf("system cannot have %d cpu(s)", n)
+       case n == 1:
+               if runtime.NumCPU() != 1 {
+                       t.Fatalf("runtime.NumCPU() returns %d on single cpu system", runtime.NumCPU())
+               }
+               return
+       }
+
+       const (
+               _CREATE_SUSPENDED   = 0x00000004
+               _PROCESS_ALL_ACCESS = syscall.STANDARD_RIGHTS_REQUIRED | syscall.SYNCHRONIZE | 0xfff
+       )
+
+       kernel32 := syscall.MustLoadDLL("kernel32.dll")
+       _GetProcessAffinityMask := kernel32.MustFindProc("GetProcessAffinityMask")
+       _SetProcessAffinityMask := kernel32.MustFindProc("SetProcessAffinityMask")
+
+       cmd := exec.Command(os.Args[0], "-test.run=TestNumCPU")
+       cmd.Env = append(os.Environ(), "GO_WANT_HELPER_PROCESS=1")
+       var buf bytes.Buffer
+       cmd.Stdout = &buf
+       cmd.Stderr = &buf
+       cmd.SysProcAttr = &syscall.SysProcAttr{CreationFlags: _CREATE_SUSPENDED}
+       err := cmd.Start()
+       if err != nil {
+               t.Fatal(err)
+       }
+       defer func() {
+               err = cmd.Wait()
+               childOutput := string(buf.Bytes())
+               if err != nil {
+                       t.Fatalf("child failed: %v: %v", err, childOutput)
+               }
+               // removeOneCPU should have decreased child cpu count by 1
+               want := fmt.Sprintf("%d", runtime.NumCPU()-1)
+               if childOutput != want {
+                       t.Fatalf("child output: want %q, got %q", want, childOutput)
+               }
+       }()
+
+       defer func() {
+               err = resumeChildThread(kernel32, cmd.Process.Pid)
+               if err != nil {
+                       t.Fatal(err)
+               }
+       }()
+
+       ph, err := syscall.OpenProcess(_PROCESS_ALL_ACCESS, false, uint32(cmd.Process.Pid))
+       if err != nil {
+               t.Fatal(err)
+       }
+       defer syscall.CloseHandle(ph)
+
+       var mask, sysmask uintptr
+       ret, _, err := _GetProcessAffinityMask.Call(uintptr(ph), uintptr(unsafe.Pointer(&mask)), uintptr(unsafe.Pointer(&sysmask)))
+       if ret == 0 {
+               t.Fatal(err)
+       }
+
+       newmask, err := removeOneCPU(mask)
+       if err != nil {
+               t.Fatal(err)
+       }
+
+       ret, _, err = _SetProcessAffinityMask.Call(uintptr(ph), newmask)
+       if ret == 0 {
+               t.Fatal(err)
+       }
+       ret, _, err = _GetProcessAffinityMask.Call(uintptr(ph), uintptr(unsafe.Pointer(&mask)), uintptr(unsafe.Pointer(&sysmask)))
+       if ret == 0 {
+               t.Fatal(err)
+       }
+       if newmask != mask {
+               t.Fatalf("SetProcessAffinityMask didn't set newmask of 0x%x. Current mask is 0x%x.", newmask, mask)
+       }
+}