gomaxprocs.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  	"fmt"
     9  	"os"
    10  	"runtime"
    11  	"strconv"
    12  	"time"
    13  )
    14  
    15  func init() {
    16  	register("PrintGOMAXPROCS", PrintGOMAXPROCS)
    17  	register("SetLimitThenDefaultGOMAXPROCS", SetLimitThenDefaultGOMAXPROCS)
    18  	register("UpdateGOMAXPROCS", UpdateGOMAXPROCS)
    19  	register("DontUpdateGOMAXPROCS", DontUpdateGOMAXPROCS)
    20  }
    21  
    22  func PrintGOMAXPROCS() {
    23  	println(runtime.GOMAXPROCS(0))
    24  }
    25  
    26  func mustSetCPUMax(path string, quota int64) {
    27  	q := "max"
    28  	if quota >= 0 {
    29  		q = strconv.FormatInt(quota, 10)
    30  	}
    31  	buf := fmt.Sprintf("%s 100000", q)
    32  	if err := os.WriteFile(path, []byte(buf), 0); err != nil {
    33  		panic(fmt.Sprintf("error setting cpu.max: %v", err))
    34  	}
    35  }
    36  
    37  func mustParseInt64(s string) int64 {
    38  	v, err := strconv.ParseInt(s, 10, 64)
    39  	if err != nil {
    40  		panic(err)
    41  	}
    42  	return v
    43  }
    44  
    45  // Inputs:
    46  // GO_TEST_CPU_MAX_PATH: Path to cgroup v2 cpu.max file.
    47  // GO_TEST_CPU_MAX_QUOTA: CPU quota to set.
    48  func SetLimitThenDefaultGOMAXPROCS() {
    49  	path := os.Getenv("GO_TEST_CPU_MAX_PATH")
    50  	quota := mustParseInt64(os.Getenv("GO_TEST_CPU_MAX_QUOTA"))
    51  
    52  	mustSetCPUMax(path, quota)
    53  
    54  	runtime.SetDefaultGOMAXPROCS()
    55  	println(runtime.GOMAXPROCS(0))
    56  }
    57  
    58  // Wait for GOMAXPROCS to change from from to to. Times out after 10s.
    59  func waitForMaxProcsChange(from, to int) {
    60  	start := time.Now()
    61  	for {
    62  		if time.Since(start) > 10*time.Second {
    63  			panic("no update for >10s")
    64  		}
    65  
    66  		procs := runtime.GOMAXPROCS(0)
    67  		println("GOMAXPROCS:", procs)
    68  		if procs == to {
    69  			return
    70  		}
    71  		if procs != from {
    72  			panic(fmt.Sprintf("GOMAXPROCS change got %d want %d", procs, to))
    73  		}
    74  
    75  		time.Sleep(100*time.Millisecond)
    76  	}
    77  }
    78  
    79  // Make sure that GOMAXPROCS does not change from curr.
    80  //
    81  // It is impossible to assert that it never changes, so this just makes sure it
    82  // stays for 5s.
    83  func mustNotChangeMaxProcs(curr int) {
    84  	start := time.Now()
    85  	for {
    86  		if time.Since(start) > 5*time.Second {
    87  			return
    88  		}
    89  
    90  		procs := runtime.GOMAXPROCS(0)
    91  		println("GOMAXPROCS:", procs)
    92  		if procs != curr {
    93  			panic(fmt.Sprintf("GOMAXPROCS change got %d want %d", procs, curr))
    94  		}
    95  
    96  		time.Sleep(100*time.Millisecond)
    97  	}
    98  }
    99  
   100  // Inputs:
   101  // GO_TEST_CPU_MAX_PATH: Path to cgroup v2 cpu.max file.
   102  func UpdateGOMAXPROCS() {
   103  	// We start with no limit.
   104  
   105  	ncpu := runtime.NumCPU()
   106  
   107  	procs := runtime.GOMAXPROCS(0)
   108  	println("GOMAXPROCS:", procs)
   109  	if procs != ncpu {
   110  		panic(fmt.Sprintf("GOMAXPROCS got %d want %d", procs, ncpu))
   111  	}
   112  
   113  	path := os.Getenv("GO_TEST_CPU_MAX_PATH")
   114  
   115  	// Drop down to 3 CPU.
   116  	mustSetCPUMax(path, 300000)
   117  	waitForMaxProcsChange(ncpu, 3)
   118  
   119  	// Drop even further. Now we hit the minimum GOMAXPROCS=2.
   120  	mustSetCPUMax(path, 100000)
   121  	waitForMaxProcsChange(3, 2)
   122  
   123  	// Increase back up.
   124  	mustSetCPUMax(path, 300000)
   125  	waitForMaxProcsChange(2, 3)
   126  
   127  	// Remove limit entirely.
   128  	mustSetCPUMax(path, -1)
   129  	waitForMaxProcsChange(3, ncpu)
   130  
   131  	// Setting GOMAXPROCS explicitly disables updates.
   132  	runtime.GOMAXPROCS(3)
   133  	mustSetCPUMax(path, 200000)
   134  	mustNotChangeMaxProcs(3)
   135  
   136  	println("OK")
   137  }
   138  
   139  // Inputs:
   140  // GO_TEST_CPU_MAX_PATH: Path to cgroup v2 cpu.max file.
   141  func DontUpdateGOMAXPROCS() {
   142  	// The caller has disabled updates. Make sure they don't happen.
   143  
   144  	curr := runtime.GOMAXPROCS(0)
   145  	println("GOMAXPROCS:", curr)
   146  
   147  	path := os.Getenv("GO_TEST_CPU_MAX_PATH")
   148  	mustSetCPUMax(path, 300000)
   149  	mustNotChangeMaxProcs(curr)
   150  
   151  	println("OK")
   152  }
   153
View as plain text