Initial QSfera import

This commit is contained in:
Курнат Андрей
2026-06-07 10:20:04 +03:00
commit 2315f25754
16485 changed files with 4826827 additions and 0 deletions
+441
View File
@@ -0,0 +1,441 @@
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 recursively.
return walkCgroupV2Hierarchy(cgroupPath, mountPoint)
}
// 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
}
// walkCgroupV2Hierarchy walks up the cgroup v2 hierarchy to find the most restrictive memory limit.
func walkCgroupV2Hierarchy(cgroupPath, mountPoint string) (uint64, error) {
var (
found = false
minLimit uint64 = math.MaxUint64
currentPath = cgroupPath
)
for {
limit, err := readMemoryLimitV2FromPath(filepath.Join(currentPath, "memory.max"))
if err != nil && !errors.Is(err, ErrNoLimit) {
return 0, err
} else if err == nil {
found = true
minLimit = min(minLimit, limit)
}
if currentPath == mountPoint {
break
}
parent := filepath.Dir(currentPath)
if parent == currentPath {
break
}
currentPath = parent
}
if !found {
return 0, ErrNoLimit
}
return minLimit, 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.stat 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.stat"))
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.stat.
// this function expects the path to be memory.stat 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.stat %q: not enough fields", line)
}
if fields[0] == "hierarchical_memory_limit" {
if len(fields) > 2 {
return 0, fmt.Errorf("failed to parse memory.stat %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.SplitN(fieldss[1], " ", 3)
if len(fields2) < 3 {
return mountInfo{}, fmt.Errorf("not enough 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
}
@@ -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
})
}
@@ -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
}
@@ -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
}
@@ -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
}
+13
View File
@@ -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 }
@@ -0,0 +1,285 @@
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)
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 spawns a goroutine that runs every refresh duration and updates the GOMEMLIMIT 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)
go func() {
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
}
}
@@ -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
}
}