gtsocial-umbx

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

memlimit.go (3153B)

      1 package memlimit
      2 
      3 import (
      4 	"errors"
      5 	"io"
      6 	"log"
      7 	"math"
      8 	"os"
      9 	"runtime/debug"
     10 	"strconv"
     11 )
     12 
     13 const (
     14 	envGOMEMLIMIT         = "GOMEMLIMIT"
     15 	envAUTOMEMLIMIT       = "AUTOMEMLIMIT"
     16 	envAUTOMEMLIMIT_DEBUG = "AUTOMEMLIMIT_DEBUG"
     17 
     18 	defaultAUTOMEMLIMIT = 0.9
     19 )
     20 
     21 var (
     22 	// ErrNoLimit is returned when the memory limit is not set.
     23 	ErrNoLimit = errors.New("memory is not limited")
     24 	// ErrNoCgroup is returned when the process is not in cgroup.
     25 	ErrNoCgroup = errors.New("process is not in cgroup")
     26 	// ErrCgroupsNotSupported is returned when the system does not support cgroups.
     27 	ErrCgroupsNotSupported = errors.New("cgroups is not supported on this system")
     28 
     29 	logger = log.New(io.Discard, "", log.LstdFlags)
     30 )
     31 
     32 // SetGoMemLimitWithEnv sets GOMEMLIMIT with the value from the environment variable.
     33 // You can configure how much memory of the cgroup's memory limit to set as GOMEMLIMIT
     34 // through AUTOMEMLIMIT in the half-open range (0.0,1.0].
     35 //
     36 // If AUTOMEMLIMIT is not set, it defaults to 0.9. (10% is the headroom for memory sources the Go runtime is unaware of.)
     37 // If GOMEMLIMIT is already set or AUTOMEMLIMIT=off, this function does nothing.
     38 func SetGoMemLimitWithEnv() {
     39 	snapshot := debug.SetMemoryLimit(-1)
     40 	defer func() {
     41 		err := recover()
     42 		if err != nil {
     43 			logger.Printf("panic during SetGoMemLimitWithEnv, rolling back to previous value %d: %v\n", snapshot, err)
     44 			debug.SetMemoryLimit(snapshot)
     45 		}
     46 	}()
     47 
     48 	if os.Getenv(envAUTOMEMLIMIT_DEBUG) == "true" {
     49 		logger = log.Default()
     50 	}
     51 
     52 	if val, ok := os.LookupEnv(envGOMEMLIMIT); ok {
     53 		logger.Printf("GOMEMLIMIT is set already, skipping: %s\n", val)
     54 		return
     55 	}
     56 
     57 	ratio := defaultAUTOMEMLIMIT
     58 	if val, ok := os.LookupEnv(envAUTOMEMLIMIT); ok {
     59 		if val == "off" {
     60 			logger.Printf("AUTOMEMLIMIT is set to off, skipping\n")
     61 			return
     62 		}
     63 		_ratio, err := strconv.ParseFloat(val, 64)
     64 		if err != nil {
     65 			logger.Printf("cannot parse AUTOMEMLIMIT: %s\n", val)
     66 			return
     67 		}
     68 		ratio = _ratio
     69 	}
     70 	if ratio <= 0 || ratio > 1 {
     71 		logger.Printf("invalid AUTOMEMLIMIT: %f\n", ratio)
     72 		return
     73 	}
     74 
     75 	limit, err := SetGoMemLimit(ratio)
     76 	if err != nil {
     77 		logger.Printf("failed to set GOMEMLIMIT: %v\n", err)
     78 		return
     79 	}
     80 
     81 	logger.Printf("GOMEMLIMIT=%d\n", limit)
     82 }
     83 
     84 // SetGoMemLimit sets GOMEMLIMIT with the value from the cgroup's memory limit and given ratio.
     85 func SetGoMemLimit(ratio float64) (int64, error) {
     86 	return SetGoMemLimitWithProvider(FromCgroup, ratio)
     87 }
     88 
     89 // Provider is a function that returns the memory limit.
     90 type Provider func() (uint64, error)
     91 
     92 // SetGoMemLimitWithProvider sets GOMEMLIMIT with the value from the given provider and ratio.
     93 func SetGoMemLimitWithProvider(provider Provider, ratio float64) (int64, error) {
     94 	limit, err := provider()
     95 	if err != nil {
     96 		return 0, err
     97 	}
     98 	goMemLimit := cappedFloat2Int(float64(limit) * ratio)
     99 	debug.SetMemoryLimit(goMemLimit)
    100 	return goMemLimit, nil
    101 }
    102 
    103 func cappedFloat2Int(f float64) int64 {
    104 	if f > math.MaxInt64 {
    105 		return math.MaxInt64
    106 	}
    107 	return int64(f)
    108 }
    109 
    110 // Limit is a helper Provider function that returns the given limit.
    111 func Limit(limit uint64) func() (uint64, error) {
    112 	return func() (uint64, error) {
    113 		return limit, nil
    114 	}
    115 }