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  	// Re-enable updates. Change is immediately visible.
   137  	runtime.SetDefaultGOMAXPROCS()
   138  	procs = runtime.GOMAXPROCS(0)
   139  	println("GOMAXPROCS:", procs)
   140  	if procs != 2 {
   141  		panic(fmt.Sprintf("GOMAXPROCS got %d want %d", procs, 2))
   142  	}
   143  
   144  	// Setting GOMAXPROCS to itself also disables updates, despite not
   145  	// changing the value itself.
   146  	runtime.GOMAXPROCS(runtime.GOMAXPROCS(0))
   147  	mustSetCPUMax(path, 300000)
   148  	mustNotChangeMaxProcs(2)
   149  
   150  	println("OK")
   151  }
   152  
   153  // Inputs:
   154  // GO_TEST_CPU_MAX_PATH: Path to cgroup v2 cpu.max file.
   155  func DontUpdateGOMAXPROCS() {
   156  	// The caller has disabled updates. Make sure they don't happen.
   157  
   158  	curr := runtime.GOMAXPROCS(0)
   159  	println("GOMAXPROCS:", curr)
   160  
   161  	path := os.Getenv("GO_TEST_CPU_MAX_PATH")
   162  	mustSetCPUMax(path, 300000)
   163  	mustNotChangeMaxProcs(curr)
   164  
   165  	println("OK")
   166  }
   167
View as plain text