diff options
Diffstat (limited to 'vendor/github.com/KimMachineGun/automemlimit')
9 files changed, 895 insertions, 0 deletions
diff --git a/vendor/github.com/KimMachineGun/automemlimit/LICENSE b/vendor/github.com/KimMachineGun/automemlimit/LICENSE new file mode 100644 index 0000000..1f5b8f6 --- /dev/null +++ b/vendor/github.com/KimMachineGun/automemlimit/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2022 Geon Kim + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/vendor/github.com/KimMachineGun/automemlimit/memlimit/cgroups.go b/vendor/github.com/KimMachineGun/automemlimit/memlimit/cgroups.go new file mode 100644 index 0000000..73a57c3 --- /dev/null +++ b/vendor/github.com/KimMachineGun/automemlimit/memlimit/cgroups.go @@ -0,0 +1,410 @@ +package memlimit + +import ( + "bufio" + "errors" + "fmt" + "io" + "math" + "os" + "path/filepath" + "slices" + "strconv" + "strings" +) + +var ( + // 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") +) + +// fromCgroup retrieves the memory limit from the cgroup. +// The versionDetector function is used to detect the cgroup version from the mountinfo. +func fromCgroup(versionDetector func(mis []mountInfo) (bool, bool)) (uint64, error) { + mf, err := os.Open("/proc/self/mountinfo") + if err != nil { + return 0, fmt.Errorf("failed to open /proc/self/mountinfo: %w", err) + } + defer mf.Close() + + mis, err := parseMountInfo(mf) + if err != nil { + return 0, fmt.Errorf("failed to parse mountinfo: %w", err) + } + + v1, v2 := versionDetector(mis) + if !(v1 || v2) { + return 0, ErrNoCgroup + } + + cf, err := os.Open("/proc/self/cgroup") + if err != nil { + return 0, fmt.Errorf("failed to open /proc/self/cgroup: %w", err) + } + defer cf.Close() + + chs, err := parseCgroupFile(cf) + if err != nil { + return 0, fmt.Errorf("failed to parse cgroup file: %w", err) + } + + if v2 { + limit, err := getMemoryLimitV2(chs, mis) + if err == nil { + return limit, nil + } else if !v1 { + return 0, err + } + } + + return getMemoryLimitV1(chs, mis) +} + +// detectCgroupVersion detects the cgroup version from the mountinfo. +func detectCgroupVersion(mis []mountInfo) (bool, bool) { + var v1, v2 bool + for _, mi := range mis { + switch mi.FilesystemType { + case "cgroup": + v1 = true + case "cgroup2": + v2 = true + } + } + return v1, v2 +} + +// getMemoryLimitV2 retrieves the memory limit from the cgroup v2 controller. +func getMemoryLimitV2(chs []cgroupHierarchy, mis []mountInfo) (uint64, error) { + // find the cgroup v2 path for the memory controller. + // in cgroup v2, the paths are unified and the controller list is empty. + idx := slices.IndexFunc(chs, func(ch cgroupHierarchy) bool { + return ch.HierarchyID == "0" && ch.ControllerList == "" + }) + if idx == -1 { + return 0, errors.New("cgroup v2 path not found") + } + relPath := chs[idx].CgroupPath + + // find the mountpoint for the cgroup v2 controller. + idx = slices.IndexFunc(mis, func(mi mountInfo) bool { + return mi.FilesystemType == "cgroup2" + }) + if idx == -1 { + return 0, errors.New("cgroup v2 mountpoint not found") + } + root, mountPoint := mis[idx].Root, mis[idx].MountPoint + + // resolve the actual cgroup path + cgroupPath, err := resolveCgroupPath(mountPoint, root, relPath) + if err != nil { + return 0, err + } + + // retrieve the memory limit from the memory.max file + return readMemoryLimitV2FromPath(filepath.Join(cgroupPath, "memory.max")) +} + +// readMemoryLimitV2FromPath reads the memory limit for cgroup v2 from the given path. +// this function expects the path to be memory.max file. +func readMemoryLimitV2FromPath(path string) (uint64, error) { + b, err := os.ReadFile(path) + if err != nil { + if errors.Is(err, os.ErrNotExist) { + return 0, ErrNoLimit + } + return 0, fmt.Errorf("failed to read memory.max: %w", err) + } + + slimit := strings.TrimSpace(string(b)) + if slimit == "max" { + return 0, ErrNoLimit + } + + limit, err := strconv.ParseUint(slimit, 10, 64) + if err != nil { + return 0, fmt.Errorf("failed to parse memory.max value: %w", err) + } + + return limit, nil +} + +// getMemoryLimitV1 retrieves the memory limit from the cgroup v1 controller. +func getMemoryLimitV1(chs []cgroupHierarchy, mis []mountInfo) (uint64, error) { + // find the cgroup v1 path for the memory controller. + idx := slices.IndexFunc(chs, func(ch cgroupHierarchy) bool { + return slices.Contains(strings.Split(ch.ControllerList, ","), "memory") + }) + if idx == -1 { + return 0, errors.New("cgroup v1 path for memory controller not found") + } + relPath := chs[idx].CgroupPath + + // find the mountpoint for the cgroup v1 controller. + idx = slices.IndexFunc(mis, func(mi mountInfo) bool { + return mi.FilesystemType == "cgroup" && slices.Contains(strings.Split(mi.SuperOptions, ","), "memory") + }) + if idx == -1 { + return 0, errors.New("cgroup v1 mountpoint for memory controller not found") + } + root, mountPoint := mis[idx].Root, mis[idx].MountPoint + + // resolve the actual cgroup path + cgroupPath, err := resolveCgroupPath(mountPoint, root, relPath) + if err != nil { + return 0, err + } + + // retrieve the memory limit from the memory.stats and memory.limit_in_bytes files. + return readMemoryLimitV1FromPath(cgroupPath) +} + +// getCgroupV1NoLimit returns the maximum value that is used to represent no limit in cgroup v1. +// the max memory limit is max int64, but it should be multiple of the page size. +func getCgroupV1NoLimit() uint64 { + ps := uint64(os.Getpagesize()) + return math.MaxInt64 / ps * ps +} + +// readMemoryLimitV1FromPath reads the memory limit for cgroup v1 from the given path. +// this function expects the path to be the cgroup directory. +func readMemoryLimitV1FromPath(cgroupPath string) (uint64, error) { + // read hierarchical_memory_limit and memory.limit_in_bytes files. + // but if hierarchical_memory_limit is not available, then use the max value as a fallback. + hml, err := readHierarchicalMemoryLimit(filepath.Join(cgroupPath, "memory.stats")) + if err != nil && !errors.Is(err, os.ErrNotExist) { + return 0, fmt.Errorf("failed to read hierarchical_memory_limit: %w", err) + } else if hml == 0 { + hml = math.MaxUint64 + } + + // read memory.limit_in_bytes file. + b, err := os.ReadFile(filepath.Join(cgroupPath, "memory.limit_in_bytes")) + if err != nil && !errors.Is(err, os.ErrNotExist) { + return 0, fmt.Errorf("failed to read memory.limit_in_bytes: %w", err) + } + lib, err := strconv.ParseUint(strings.TrimSpace(string(b)), 10, 64) + if err != nil { + return 0, fmt.Errorf("failed to parse memory.limit_in_bytes value: %w", err) + } else if lib == 0 { + hml = math.MaxUint64 + } + + // use the minimum value between hierarchical_memory_limit and memory.limit_in_bytes. + // if the limit is the maximum value, then it is considered as no limit. + limit := min(hml, lib) + if limit >= getCgroupV1NoLimit() { + return 0, ErrNoLimit + } + + return limit, nil +} + +// readHierarchicalMemoryLimit extracts hierarchical_memory_limit from memory.stats. +// this function expects the path to be memory.stats file. +func readHierarchicalMemoryLimit(path string) (uint64, error) { + file, err := os.Open(path) + if err != nil { + return 0, err + } + defer file.Close() + + scanner := bufio.NewScanner(file) + for scanner.Scan() { + line := scanner.Text() + + fields := strings.Split(line, " ") + if len(fields) < 2 { + return 0, fmt.Errorf("failed to parse memory.stats %q: not enough fields", line) + } + + if fields[0] == "hierarchical_memory_limit" { + if len(fields) > 2 { + return 0, fmt.Errorf("failed to parse memory.stats %q: too many fields for hierarchical_memory_limit", line) + } + return strconv.ParseUint(fields[1], 10, 64) + } + } + if err := scanner.Err(); err != nil { + return 0, err + } + + return 0, nil +} + +// https://www.man7.org/linux/man-pages/man5/proc_pid_mountinfo.5.html +// 731 771 0:59 /sysrq-trigger /proc/sysrq-trigger ro,nosuid,nodev,noexec,relatime - proc proc rw +// +// 36 35 98:0 /mnt1 /mnt2 rw,noatime master:1 - ext3 /dev/root rw,errors=continue +// (1)(2)(3) (4) (5) (6) (7) (8) (9) (10) (11) +// +// (1) mount ID: a unique ID for the mount (may be reused after umount(2)). +// (2) parent ID: the ID of the parent mount (or of self for the root of this mount namespace's mount tree). +// (3) major:minor: the value of st_dev for files on this filesystem (see stat(2)). +// (4) root: the pathname of the directory in the filesystem which forms the root of this mount. +// (5) mount point: the pathname of the mount point relative to the process's root directory. +// (6) mount options: per-mount options (see mount(2)). +// (7) optional fields: zero or more fields of the form "tag[:value]"; see below. +// (8) separator: the end of the optional fields is marked by a single hyphen. +// (9) filesystem type: the filesystem type in the form "type[.subtype]". +// (10) mount source: filesystem-specific information or "none". +// (11) super options: per-superblock options (see mount(2)). +type mountInfo struct { + Root string + MountPoint string + FilesystemType string + SuperOptions string +} + +// parseMountInfoLine parses a line from the mountinfo file. +func parseMountInfoLine(line string) (mountInfo, error) { + if line == "" { + return mountInfo{}, errors.New("empty line") + } + + fieldss := strings.SplitN(line, " - ", 2) + if len(fieldss) != 2 { + return mountInfo{}, fmt.Errorf("invalid separator") + } + + fields1 := strings.SplitN(fieldss[0], " ", 7) + if len(fields1) < 6 { + return mountInfo{}, fmt.Errorf("not enough fields before separator: %v", fields1) + } else if len(fields1) == 6 { + fields1 = append(fields1, "") + } + + fields2 := strings.Split(fieldss[1], " ") + if len(fields2) < 3 { + return mountInfo{}, fmt.Errorf("not enough fields after separator: %v", fields2) + } else if len(fields2) > 3 { + return mountInfo{}, fmt.Errorf("too many fields after separator: %v", fields2) + } + + return mountInfo{ + Root: fields1[3], + MountPoint: fields1[4], + FilesystemType: fields2[0], + SuperOptions: fields2[2], + }, nil +} + +// parseMountInfo parses the mountinfo file. +func parseMountInfo(r io.Reader) ([]mountInfo, error) { + var ( + s = bufio.NewScanner(r) + mis []mountInfo + ) + for s.Scan() { + line := s.Text() + + mi, err := parseMountInfoLine(line) + if err != nil { + return nil, fmt.Errorf("failed to parse mountinfo file %q: %w", line, err) + } + + mis = append(mis, mi) + } + if err := s.Err(); err != nil { + return nil, err + } + + return mis, nil +} + +// https://www.man7.org/linux/man-pages/man7/cgroups.7.html +// +// 5:cpuacct,cpu,cpuset:/daemons +// (1) (2) (3) +// +// (1) hierarchy ID: +// +// cgroups version 1 hierarchies, this field +// contains a unique hierarchy ID number that can be +// matched to a hierarchy ID in /proc/cgroups. For the +// cgroups version 2 hierarchy, this field contains the +// value 0. +// +// (2) controller list: +// +// For cgroups version 1 hierarchies, this field +// contains a comma-separated list of the controllers +// bound to the hierarchy. For the cgroups version 2 +// hierarchy, this field is empty. +// +// (3) cgroup path: +// +// This field contains the pathname of the control group +// in the hierarchy to which the process belongs. This +// pathname is relative to the mount point of the +// hierarchy. +type cgroupHierarchy struct { + HierarchyID string + ControllerList string + CgroupPath string +} + +// parseCgroupHierarchyLine parses a line from the cgroup file. +func parseCgroupHierarchyLine(line string) (cgroupHierarchy, error) { + if line == "" { + return cgroupHierarchy{}, errors.New("empty line") + } + + fields := strings.Split(line, ":") + if len(fields) < 3 { + return cgroupHierarchy{}, fmt.Errorf("not enough fields: %v", fields) + } else if len(fields) > 3 { + return cgroupHierarchy{}, fmt.Errorf("too many fields: %v", fields) + } + + return cgroupHierarchy{ + HierarchyID: fields[0], + ControllerList: fields[1], + CgroupPath: fields[2], + }, nil +} + +// parseCgroupFile parses the cgroup file. +func parseCgroupFile(r io.Reader) ([]cgroupHierarchy, error) { + var ( + s = bufio.NewScanner(r) + chs []cgroupHierarchy + ) + for s.Scan() { + line := s.Text() + + ch, err := parseCgroupHierarchyLine(line) + if err != nil { + return nil, fmt.Errorf("failed to parse cgroup file %q: %w", line, err) + } + + chs = append(chs, ch) + } + if err := s.Err(); err != nil { + return nil, err + } + + return chs, nil +} + +// resolveCgroupPath resolves the actual cgroup path from the mountpoint, root, and cgroupRelPath. +func resolveCgroupPath(mountpoint, root, cgroupRelPath string) (string, error) { + rel, err := filepath.Rel(root, cgroupRelPath) + if err != nil { + return "", err + } + + // if the relative path is ".", then the cgroupRelPath is the root itself. + if rel == "." { + return mountpoint, nil + } + + // if the relative path starts with "..", then it is outside the root. + if strings.HasPrefix(rel, "..") { + return "", fmt.Errorf("invalid cgroup path: %s is not under root %s", cgroupRelPath, root) + } + + return filepath.Join(mountpoint, rel), nil +} diff --git a/vendor/github.com/KimMachineGun/automemlimit/memlimit/cgroups_linux.go b/vendor/github.com/KimMachineGun/automemlimit/memlimit/cgroups_linux.go new file mode 100644 index 0000000..fd2c7e4 --- /dev/null +++ b/vendor/github.com/KimMachineGun/automemlimit/memlimit/cgroups_linux.go @@ -0,0 +1,32 @@ +//go:build linux +// +build linux + +package memlimit + +// FromCgroup retrieves the memory limit from the cgroup. +func FromCgroup() (uint64, error) { + return fromCgroup(detectCgroupVersion) +} + +// FromCgroupV1 retrieves the memory limit from the cgroup v1 controller. +// After v1.0.0, this function could be removed and FromCgroup should be used instead. +func FromCgroupV1() (uint64, error) { + return fromCgroup(func(_ []mountInfo) (bool, bool) { + return true, false + }) +} + +// FromCgroupHybrid retrieves the memory limit from the cgroup v2 and v1 controller sequentially, +// basically, it is equivalent to FromCgroup. +// After v1.0.0, this function could be removed and FromCgroup should be used instead. +func FromCgroupHybrid() (uint64, error) { + return FromCgroup() +} + +// FromCgroupV2 retrieves the memory limit from the cgroup v2 controller. +// After v1.0.0, this function could be removed and FromCgroup should be used instead. +func FromCgroupV2() (uint64, error) { + return fromCgroup(func(_ []mountInfo) (bool, bool) { + return false, true + }) +} diff --git a/vendor/github.com/KimMachineGun/automemlimit/memlimit/cgroups_unsupported.go b/vendor/github.com/KimMachineGun/automemlimit/memlimit/cgroups_unsupported.go new file mode 100644 index 0000000..9feca81 --- /dev/null +++ b/vendor/github.com/KimMachineGun/automemlimit/memlimit/cgroups_unsupported.go @@ -0,0 +1,20 @@ +//go:build !linux +// +build !linux + +package memlimit + +func FromCgroup() (uint64, error) { + return 0, ErrCgroupsNotSupported +} + +func FromCgroupV1() (uint64, error) { + return 0, ErrCgroupsNotSupported +} + +func FromCgroupHybrid() (uint64, error) { + return 0, ErrCgroupsNotSupported +} + +func FromCgroupV2() (uint64, error) { + return 0, ErrCgroupsNotSupported +} diff --git a/vendor/github.com/KimMachineGun/automemlimit/memlimit/exp_system.go b/vendor/github.com/KimMachineGun/automemlimit/memlimit/exp_system.go new file mode 100644 index 0000000..dee95f5 --- /dev/null +++ b/vendor/github.com/KimMachineGun/automemlimit/memlimit/exp_system.go @@ -0,0 +1,14 @@ +package memlimit + +import ( + "github.com/pbnjay/memory" +) + +// FromSystem returns the total memory of the system. +func FromSystem() (uint64, error) { + limit := memory.TotalMemory() + if limit == 0 { + return 0, ErrNoLimit + } + return limit, nil +} diff --git a/vendor/github.com/KimMachineGun/automemlimit/memlimit/experiment.go b/vendor/github.com/KimMachineGun/automemlimit/memlimit/experiment.go new file mode 100644 index 0000000..2a7c320 --- /dev/null +++ b/vendor/github.com/KimMachineGun/automemlimit/memlimit/experiment.go @@ -0,0 +1,59 @@ +package memlimit + +import ( + "fmt" + "os" + "reflect" + "strings" +) + +const ( + envAUTOMEMLIMIT_EXPERIMENT = "AUTOMEMLIMIT_EXPERIMENT" +) + +// Experiments is a set of experiment flags. +// It is used to enable experimental features. +// +// You can set the flags by setting the environment variable AUTOMEMLIMIT_EXPERIMENT. +// The value of the environment variable is a comma-separated list of experiment names. +// +// The following experiment names are known: +// +// - none: disable all experiments +// - system: enable fallback to system memory limit +type Experiments struct { + // System enables fallback to system memory limit. + System bool +} + +func parseExperiments() (Experiments, error) { + var exp Experiments + + // Create a map of known experiment names. + names := make(map[string]func(bool)) + rv := reflect.ValueOf(&exp).Elem() + rt := rv.Type() + for i := 0; i < rt.NumField(); i++ { + field := rv.Field(i) + names[strings.ToLower(rt.Field(i).Name)] = field.SetBool + } + + // Parse names. + for _, f := range strings.Split(os.Getenv(envAUTOMEMLIMIT_EXPERIMENT), ",") { + if f == "" { + continue + } + if f == "none" { + exp = Experiments{} + continue + } + val := true + set, ok := names[f] + if !ok { + return Experiments{}, fmt.Errorf("unknown AUTOMEMLIMIT_EXPERIMENT %s", f) + } + set(val) + } + + return exp, nil +} diff --git a/vendor/github.com/KimMachineGun/automemlimit/memlimit/logger.go b/vendor/github.com/KimMachineGun/automemlimit/memlimit/logger.go new file mode 100644 index 0000000..4cf0b58 --- /dev/null +++ b/vendor/github.com/KimMachineGun/automemlimit/memlimit/logger.go @@ -0,0 +1,13 @@ +package memlimit + +import ( + "context" + "log/slog" +) + +type noopLogger struct{} + +func (noopLogger) Enabled(context.Context, slog.Level) bool { return false } +func (noopLogger) Handle(context.Context, slog.Record) error { return nil } +func (d noopLogger) WithAttrs([]slog.Attr) slog.Handler { return d } +func (d noopLogger) WithGroup(string) slog.Handler { return d } diff --git a/vendor/github.com/KimMachineGun/automemlimit/memlimit/memlimit.go b/vendor/github.com/KimMachineGun/automemlimit/memlimit/memlimit.go new file mode 100644 index 0000000..cbd53ce --- /dev/null +++ b/vendor/github.com/KimMachineGun/automemlimit/memlimit/memlimit.go @@ -0,0 +1,283 @@ +package memlimit + +import ( + "errors" + "fmt" + "log/slog" + "math" + "os" + "runtime/debug" + "strconv" + "time" +) + +const ( + envGOMEMLIMIT = "GOMEMLIMIT" + envAUTOMEMLIMIT = "AUTOMEMLIMIT" + // Deprecated: use memlimit.WithLogger instead + envAUTOMEMLIMIT_DEBUG = "AUTOMEMLIMIT_DEBUG" + + defaultAUTOMEMLIMIT = 0.9 +) + +// ErrNoLimit is returned when the memory limit is not set. +var ErrNoLimit = errors.New("memory is not limited") + +type config struct { + logger *slog.Logger + ratio float64 + provider Provider + refresh time.Duration +} + +// Option is a function that configures the behavior of SetGoMemLimitWithOptions. +type Option func(cfg *config) + +// WithRatio configures the ratio of the memory limit to set as GOMEMLIMIT. +// +// Default: 0.9 +func WithRatio(ratio float64) Option { + return func(cfg *config) { + cfg.ratio = ratio + } +} + +// WithProvider configures the provider. +// +// Default: FromCgroup +func WithProvider(provider Provider) Option { + return func(cfg *config) { + cfg.provider = provider + } +} + +// WithLogger configures the logger. +// It automatically attaches the "package" attribute to the logs. +// +// Default: slog.New(noopLogger{}) +func WithLogger(logger *slog.Logger) Option { + return func(cfg *config) { + cfg.logger = memlimitLogger(logger) + } +} + +// WithRefreshInterval configures the refresh interval for automemlimit. +// If a refresh interval is greater than 0, automemlimit periodically fetches +// the memory limit from the provider and reapplies it if it has changed. +// If the provider returns an error, it logs the error and continues. +// ErrNoLimit is treated as math.MaxInt64. +// +// Default: 0 (no refresh) +func WithRefreshInterval(refresh time.Duration) Option { + return func(cfg *config) { + cfg.refresh = refresh + } +} + +// WithEnv configures whether to use environment variables. +// +// Default: false +// +// Deprecated: currently this does nothing. +func WithEnv() Option { + return func(cfg *config) {} +} + +func memlimitLogger(logger *slog.Logger) *slog.Logger { + if logger == nil { + return slog.New(noopLogger{}) + } + return logger.With(slog.String("package", "github.com/KimMachineGun/automemlimit/memlimit")) +} + +// SetGoMemLimitWithOpts sets GOMEMLIMIT with options and environment variables. +// +// You can configure how much memory of the cgroup's memory limit to set as GOMEMLIMIT +// through AUTOMEMLIMIT environment variable 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. +// +// If AUTOMEMLIMIT_EXPERIMENT is set, it enables experimental features. +// Please see the documentation of Experiments for more details. +// +// Options: +// - WithRatio +// - WithProvider +// - WithLogger +func SetGoMemLimitWithOpts(opts ...Option) (_ int64, _err error) { + // init config + cfg := &config{ + logger: slog.New(noopLogger{}), + ratio: defaultAUTOMEMLIMIT, + provider: FromCgroup, + } + // TODO: remove this + if debug, ok := os.LookupEnv(envAUTOMEMLIMIT_DEBUG); ok { + defaultLogger := memlimitLogger(slog.Default()) + defaultLogger.Warn("AUTOMEMLIMIT_DEBUG is deprecated, use memlimit.WithLogger instead") + if debug == "true" { + cfg.logger = defaultLogger + } + } + for _, opt := range opts { + opt(cfg) + } + + // log error if any on return + defer func() { + if _err != nil { + cfg.logger.Error("failed to set GOMEMLIMIT", slog.Any("error", _err)) + } + }() + + // parse experiments + exps, err := parseExperiments() + if err != nil { + return 0, fmt.Errorf("failed to parse experiments: %w", err) + } + if exps.System { + cfg.logger.Info("system experiment is enabled: using system memory limit as a fallback") + cfg.provider = ApplyFallback(cfg.provider, FromSystem) + } + + // rollback to previous memory limit on panic + snapshot := debug.SetMemoryLimit(-1) + defer rollbackOnPanic(cfg.logger, snapshot, &_err) + + // check if GOMEMLIMIT is already set + if val, ok := os.LookupEnv(envGOMEMLIMIT); ok { + cfg.logger.Info("GOMEMLIMIT is already set, skipping", slog.String(envGOMEMLIMIT, val)) + return 0, nil + } + + // parse AUTOMEMLIMIT + ratio := cfg.ratio + if val, ok := os.LookupEnv(envAUTOMEMLIMIT); ok { + if val == "off" { + cfg.logger.Info("AUTOMEMLIMIT is set to off, skipping") + return 0, nil + } + ratio, err = strconv.ParseFloat(val, 64) + if err != nil { + return 0, fmt.Errorf("cannot parse AUTOMEMLIMIT: %s", val) + } + } + + // apply ratio to the provider + provider := capProvider(ApplyRatio(cfg.provider, ratio)) + + // set the memory limit and start refresh + limit, err := updateGoMemLimit(uint64(snapshot), provider, cfg.logger) + go refresh(provider, cfg.logger, cfg.refresh) + if err != nil { + if errors.Is(err, ErrNoLimit) { + cfg.logger.Info("memory is not limited, skipping") + // TODO: consider returning the snapshot + return 0, nil + } + return 0, fmt.Errorf("failed to set GOMEMLIMIT: %w", err) + } + + return int64(limit), nil +} + +// updateGoMemLimit updates the Go's memory limit, if it has changed. +func updateGoMemLimit(currLimit uint64, provider Provider, logger *slog.Logger) (uint64, error) { + newLimit, err := provider() + if err != nil { + return 0, err + } + + if newLimit == currLimit { + logger.Debug("GOMEMLIMIT is not changed, skipping", slog.Uint64(envGOMEMLIMIT, newLimit)) + return newLimit, nil + } + + debug.SetMemoryLimit(int64(newLimit)) + logger.Info("GOMEMLIMIT is updated", slog.Uint64(envGOMEMLIMIT, newLimit), slog.Uint64("previous", currLimit)) + + return newLimit, nil +} + +// refresh periodically fetches the memory limit from the provider and reapplies it if it has changed. +// See more details in the documentation of WithRefreshInterval. +func refresh(provider Provider, logger *slog.Logger, refresh time.Duration) { + if refresh == 0 { + return + } + + provider = noErrNoLimitProvider(provider) + + t := time.NewTicker(refresh) + for range t.C { + err := func() (_err error) { + snapshot := debug.SetMemoryLimit(-1) + defer rollbackOnPanic(logger, snapshot, &_err) + + _, err := updateGoMemLimit(uint64(snapshot), provider, logger) + if err != nil { + return err + } + + return nil + }() + if err != nil { + logger.Error("failed to refresh GOMEMLIMIT", slog.Any("error", err)) + } + } +} + +// rollbackOnPanic rollbacks to the snapshot on panic. +// Since it uses recover, it should be called in a deferred function. +func rollbackOnPanic(logger *slog.Logger, snapshot int64, err *error) { + panicErr := recover() + if panicErr != nil { + if *err != nil { + logger.Error("failed to set GOMEMLIMIT", slog.Any("error", *err)) + } + *err = fmt.Errorf("panic during setting the Go's memory limit, rolling back to previous limit %d: %v", + snapshot, panicErr, + ) + debug.SetMemoryLimit(snapshot) + } +} + +// SetGoMemLimitWithEnv sets GOMEMLIMIT with the value from the environment variables. +// Since WithEnv is deprecated, this function is equivalent to SetGoMemLimitWithOpts(). +// Deprecated: use SetGoMemLimitWithOpts instead. +func SetGoMemLimitWithEnv() { + _, _ = SetGoMemLimitWithOpts() +} + +// SetGoMemLimit sets GOMEMLIMIT with the value from the cgroup's memory limit and given ratio. +func SetGoMemLimit(ratio float64) (int64, error) { + return SetGoMemLimitWithOpts(WithRatio(ratio)) +} + +// SetGoMemLimitWithProvider sets GOMEMLIMIT with the value from the given provider and ratio. +func SetGoMemLimitWithProvider(provider Provider, ratio float64) (int64, error) { + return SetGoMemLimitWithOpts(WithProvider(provider), WithRatio(ratio)) +} + +func noErrNoLimitProvider(provider Provider) Provider { + return func() (uint64, error) { + limit, err := provider() + if errors.Is(err, ErrNoLimit) { + return math.MaxInt64, nil + } + return limit, err + } +} + +func capProvider(provider Provider) Provider { + return func() (uint64, error) { + limit, err := provider() + if err != nil { + return 0, err + } else if limit > math.MaxInt64 { + return math.MaxInt64, nil + } + return limit, nil + } +} diff --git a/vendor/github.com/KimMachineGun/automemlimit/memlimit/provider.go b/vendor/github.com/KimMachineGun/automemlimit/memlimit/provider.go new file mode 100644 index 0000000..4f83770 --- /dev/null +++ b/vendor/github.com/KimMachineGun/automemlimit/memlimit/provider.go @@ -0,0 +1,43 @@ +package memlimit + +import ( + "fmt" +) + +// Provider is a function that returns the memory limit. +type Provider func() (uint64, error) + +// 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 + } +} + +// ApplyRationA is a helper Provider function that applies the given ratio to the given provider. +func ApplyRatio(provider Provider, ratio float64) Provider { + if ratio == 1 { + return provider + } + return func() (uint64, error) { + if ratio <= 0 || ratio > 1 { + return 0, fmt.Errorf("invalid ratio: %f, ratio should be in the range (0.0,1.0]", ratio) + } + limit, err := provider() + if err != nil { + return 0, err + } + return uint64(float64(limit) * ratio), nil + } +} + +// ApplyFallback is a helper Provider function that sets the fallback provider. +func ApplyFallback(provider Provider, fallback Provider) Provider { + return func() (uint64, error) { + limit, err := provider() + if err != nil { + return fallback() + } + return limit, nil + } +} |