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