gomaxprocs_windows.go

     1  // Copyright 2025 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package main
     6  
     7  import (
     8  	"os"
     9  	"runtime"
    10  	"syscall"
    11  	"unsafe"
    12  )
    13  
    14  func init() {
    15  	register("WindowsUpdateGOMAXPROCS", WindowsUpdateGOMAXPROCS)
    16  	register("WindowsDontUpdateGOMAXPROCS", WindowsDontUpdateGOMAXPROCS)
    17  }
    18  
    19  // Set CPU affinity mask to only two CPUs.
    20  //
    21  // Skips the test if CPUs 0 and 1 are not available.
    22  func setAffinity2() {
    23  	kernel32 := syscall.MustLoadDLL("kernel32.dll")
    24  	_GetProcessAffinityMask := kernel32.MustFindProc("GetProcessAffinityMask")
    25  	_SetProcessAffinityMask := kernel32.MustFindProc("SetProcessAffinityMask")
    26  
    27  	h, err := syscall.GetCurrentProcess()
    28  	if err != nil {
    29  		panic(err)
    30  	}
    31  
    32  	var mask, sysmask uintptr
    33  	ret, _, err := _GetProcessAffinityMask.Call(uintptr(h), uintptr(unsafe.Pointer(&mask)), uintptr(unsafe.Pointer(&sysmask)))
    34  	if ret == 0 {
    35  		panic(err)
    36  	}
    37  
    38  	// We're going to restrict to CPUs 0 and 1. Make sure those are already available.
    39  	if mask & 0b11 != 0b11 {
    40  		println("SKIP: CPUs 0 and 1 not available")
    41  		os.Exit(0)
    42  	}
    43  
    44  	mask = 0b11
    45  	ret, _, err = _SetProcessAffinityMask.Call(uintptr(h), mask)
    46  	if ret == 0 {
    47  		panic(err)
    48  	}
    49  }
    50  
    51  func WindowsUpdateGOMAXPROCS() {
    52  	ncpu := runtime.NumCPU()
    53  	setAffinity2()
    54  	waitForMaxProcsChange(ncpu, 2)
    55  	println("OK")
    56  }
    57  
    58  func WindowsDontUpdateGOMAXPROCS() {
    59  	procs := runtime.GOMAXPROCS(0)
    60  	setAffinity2()
    61  	mustNotChangeMaxProcs(procs)
    62  	println("OK")
    63  }
    64
View as plain text