gtsocial-umbx

Unnamed repository; edit this file 'description' to name the repository.
Log | Files | Refs | README | LICENSE

maxprocs.go (4055B)

      1 // Copyright (c) 2017 Uber Technologies, Inc.
      2 //
      3 // Permission is hereby granted, free of charge, to any person obtaining a copy
      4 // of this software and associated documentation files (the "Software"), to deal
      5 // in the Software without restriction, including without limitation the rights
      6 // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
      7 // copies of the Software, and to permit persons to whom the Software is
      8 // furnished to do so, subject to the following conditions:
      9 //
     10 // The above copyright notice and this permission notice shall be included in
     11 // all copies or substantial portions of the Software.
     12 //
     13 // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
     14 // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
     15 // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
     16 // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
     17 // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
     18 // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
     19 // THE SOFTWARE.
     20 
     21 // Package maxprocs lets Go programs easily configure runtime.GOMAXPROCS to
     22 // match the configured Linux CPU quota. Unlike the top-level automaxprocs
     23 // package, it lets the caller configure logging and handle errors.
     24 package maxprocs // import "go.uber.org/automaxprocs/maxprocs"
     25 
     26 import (
     27 	"os"
     28 	"runtime"
     29 
     30 	iruntime "go.uber.org/automaxprocs/internal/runtime"
     31 )
     32 
     33 const _maxProcsKey = "GOMAXPROCS"
     34 
     35 func currentMaxProcs() int {
     36 	return runtime.GOMAXPROCS(0)
     37 }
     38 
     39 type config struct {
     40 	printf        func(string, ...interface{})
     41 	procs         func(int) (int, iruntime.CPUQuotaStatus, error)
     42 	minGOMAXPROCS int
     43 }
     44 
     45 func (c *config) log(fmt string, args ...interface{}) {
     46 	if c.printf != nil {
     47 		c.printf(fmt, args...)
     48 	}
     49 }
     50 
     51 // An Option alters the behavior of Set.
     52 type Option interface {
     53 	apply(*config)
     54 }
     55 
     56 // Logger uses the supplied printf implementation for log output. By default,
     57 // Set doesn't log anything.
     58 func Logger(printf func(string, ...interface{})) Option {
     59 	return optionFunc(func(cfg *config) {
     60 		cfg.printf = printf
     61 	})
     62 }
     63 
     64 // Min sets the minimum GOMAXPROCS value that will be used.
     65 // Any value below 1 is ignored.
     66 func Min(n int) Option {
     67 	return optionFunc(func(cfg *config) {
     68 		if n >= 1 {
     69 			cfg.minGOMAXPROCS = n
     70 		}
     71 	})
     72 }
     73 
     74 type optionFunc func(*config)
     75 
     76 func (of optionFunc) apply(cfg *config) { of(cfg) }
     77 
     78 // Set GOMAXPROCS to match the Linux container CPU quota (if any), returning
     79 // any error encountered and an undo function.
     80 //
     81 // Set is a no-op on non-Linux systems and in Linux environments without a
     82 // configured CPU quota.
     83 func Set(opts ...Option) (func(), error) {
     84 	cfg := &config{
     85 		procs:         iruntime.CPUQuotaToGOMAXPROCS,
     86 		minGOMAXPROCS: 1,
     87 	}
     88 	for _, o := range opts {
     89 		o.apply(cfg)
     90 	}
     91 
     92 	undoNoop := func() {
     93 		cfg.log("maxprocs: No GOMAXPROCS change to reset")
     94 	}
     95 
     96 	// Honor the GOMAXPROCS environment variable if present. Otherwise, amend
     97 	// `runtime.GOMAXPROCS()` with the current process' CPU quota if the OS is
     98 	// Linux, and guarantee a minimum value of 1. The minimum guaranteed value
     99 	// can be overridden using `maxprocs.Min()`.
    100 	if max, exists := os.LookupEnv(_maxProcsKey); exists {
    101 		cfg.log("maxprocs: Honoring GOMAXPROCS=%q as set in environment", max)
    102 		return undoNoop, nil
    103 	}
    104 
    105 	maxProcs, status, err := cfg.procs(cfg.minGOMAXPROCS)
    106 	if err != nil {
    107 		return undoNoop, err
    108 	}
    109 
    110 	if status == iruntime.CPUQuotaUndefined {
    111 		cfg.log("maxprocs: Leaving GOMAXPROCS=%v: CPU quota undefined", currentMaxProcs())
    112 		return undoNoop, nil
    113 	}
    114 
    115 	prev := currentMaxProcs()
    116 	undo := func() {
    117 		cfg.log("maxprocs: Resetting GOMAXPROCS to %v", prev)
    118 		runtime.GOMAXPROCS(prev)
    119 	}
    120 
    121 	switch status {
    122 	case iruntime.CPUQuotaMinUsed:
    123 		cfg.log("maxprocs: Updating GOMAXPROCS=%v: using minimum allowed GOMAXPROCS", maxProcs)
    124 	case iruntime.CPUQuotaUsed:
    125 		cfg.log("maxprocs: Updating GOMAXPROCS=%v: determined from CPU quota", maxProcs)
    126 	}
    127 
    128 	runtime.GOMAXPROCS(maxProcs)
    129 	return undo, nil
    130 }