GoToSocial/vendor/github.com/KimMachineGun/automemlimit/memlimit/memlimit.go

214 lines
5.6 KiB
Go
Raw Normal View History

package memlimit
import (
"errors"
"fmt"
"log/slog"
"math"
"os"
"runtime/debug"
"strconv"
)
const (
envGOMEMLIMIT = "GOMEMLIMIT"
envAUTOMEMLIMIT = "AUTOMEMLIMIT"
// Deprecated: use memlimit.WithLogger instead
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")
)
type config struct {
logger *slog.Logger
ratio float64
provider Provider
}
// 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)
}
}
// 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 envrironment 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)
}
// capture the current GOMEMLIMIT for rollback in case of panic
snapshot := debug.SetMemoryLimit(-1)
defer func() {
panicErr := recover()
if panicErr != nil {
if _err != nil {
cfg.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)
}
}()
// 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)
}
ratio = _ratio
}
// set GOMEMLIMIT
limit, err := setGoMemLimit(ApplyRatio(cfg.provider, ratio))
if err != nil {
if errors.Is(err, ErrNoLimit) {
[chore] dependabot updates (#2922) * [chore]: Bump github.com/prometheus/client_golang from 1.18.0 to 1.19.1 Bumps [github.com/prometheus/client_golang](https://github.com/prometheus/client_golang) from 1.18.0 to 1.19.1. - [Release notes](https://github.com/prometheus/client_golang/releases) - [Changelog](https://github.com/prometheus/client_golang/blob/main/CHANGELOG.md) - [Commits](https://github.com/prometheus/client_golang/compare/v1.18.0...v1.19.1) --- updated-dependencies: - dependency-name: github.com/prometheus/client_golang dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] <support@github.com> * [chore]: Bump github.com/KimMachineGun/automemlimit from 0.6.0 to 0.6.1 Bumps [github.com/KimMachineGun/automemlimit](https://github.com/KimMachineGun/automemlimit) from 0.6.0 to 0.6.1. - [Release notes](https://github.com/KimMachineGun/automemlimit/releases) - [Commits](https://github.com/KimMachineGun/automemlimit/compare/v0.6.0...v0.6.1) --- updated-dependencies: - dependency-name: github.com/KimMachineGun/automemlimit dependency-type: direct:production update-type: version-update:semver-patch ... Signed-off-by: dependabot[bot] <support@github.com> * [chore]: Bump github.com/tdewolff/minify/v2 from 2.20.20 to 2.20.24 Bumps [github.com/tdewolff/minify/v2](https://github.com/tdewolff/minify) from 2.20.20 to 2.20.24. - [Release notes](https://github.com/tdewolff/minify/releases) - [Commits](https://github.com/tdewolff/minify/compare/v2.20.20...v2.20.24) --- updated-dependencies: - dependency-name: github.com/tdewolff/minify/v2 dependency-type: direct:production update-type: version-update:semver-patch ... Signed-off-by: dependabot[bot] <support@github.com> * [chore]: Bump github.com/go-swagger/go-swagger Bumps [github.com/go-swagger/go-swagger](https://github.com/go-swagger/go-swagger) from 0.30.6-0.20240418033037-c46c303aaa02 to 0.31.0. - [Release notes](https://github.com/go-swagger/go-swagger/releases) - [Changelog](https://github.com/go-swagger/go-swagger/blob/master/.goreleaser.yml) - [Commits](https://github.com/go-swagger/go-swagger/commits/v0.31.0) --- updated-dependencies: - dependency-name: github.com/go-swagger/go-swagger dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] <support@github.com> * [chore]: Bump github.com/gin-gonic/gin from 1.9.1 to 1.10.0 Bumps [github.com/gin-gonic/gin](https://github.com/gin-gonic/gin) from 1.9.1 to 1.10.0. - [Release notes](https://github.com/gin-gonic/gin/releases) - [Changelog](https://github.com/gin-gonic/gin/blob/master/CHANGELOG.md) - [Commits](https://github.com/gin-gonic/gin/compare/v1.9.1...v1.10.0) --- updated-dependencies: - dependency-name: github.com/gin-gonic/gin dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] <support@github.com> --------- Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
2024-05-21 15:17:22 +02:00
cfg.logger.Info("memory is not limited, skipping")
return 0, nil
}
return 0, fmt.Errorf("failed to set GOMEMLIMIT: %w", err)
}
cfg.logger.Info("GOMEMLIMIT is updated", slog.Int64(envGOMEMLIMIT, limit))
return limit, nil
}
// 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 setGoMemLimit(provider Provider) (int64, error) {
limit, err := provider()
if err != nil {
return 0, err
}
capped := cappedU64ToI64(limit)
debug.SetMemoryLimit(capped)
return capped, nil
}
func cappedU64ToI64(limit uint64) int64 {
if limit > math.MaxInt64 {
return math.MaxInt64
}
return int64(limit)
}