pax_global_header00006660000000000000000000000064147565552450014534gustar00rootroot0000000000000052 comment=e027e7ec771daa7242ca1dfee321b77e0ead763e golang-github-kimmachinegun-automemlimit-0.7.1/000077500000000000000000000000001475655524500216115ustar00rootroot00000000000000golang-github-kimmachinegun-automemlimit-0.7.1/.github/000077500000000000000000000000001475655524500231515ustar00rootroot00000000000000golang-github-kimmachinegun-automemlimit-0.7.1/.github/workflows/000077500000000000000000000000001475655524500252065ustar00rootroot00000000000000golang-github-kimmachinegun-automemlimit-0.7.1/.github/workflows/test.yml000066400000000000000000000033651475655524500267170ustar00rootroot00000000000000name: Test on: [ push, pull_request ] jobs: test-ubuntu-20_04: runs-on: ubuntu-20.04 steps: - uses: actions/checkout@v3 - name: Docker Info run: | docker info - name: Pull golang image run: | docker pull golang:1.22 - name: Run tests in Go container (1000m) run: | docker run --rm -v=$(pwd):/app -w=/app -m=1000m golang:1.22 go test -v ./... -expected=1048576000 -cgroup-version 1 - name: Run tests in Go container (4321m) run: | docker run --rm -v=$(pwd):/app -w=/app -m=4321m golang:1.22 go test -v ./... -expected=4530896896 -cgroup-version 1 - name: Run tests in Go container (system memory limit) run: | docker run --rm -v=$(pwd):/app -w=/app golang:1.22 go test -v ./... -expected-system=$(($(awk '/MemTotal/ {print $2}' /proc/meminfo) * 1024)) -cgroup-version 1 test-ubuntu-22_04: runs-on: ubuntu-22.04 steps: - uses: actions/checkout@v3 - name: Docker Info run: | docker info - name: Pull golang image run: | docker pull golang:1.22 - name: Run tests in Go container (1000m) run: | docker run --rm -v=$(pwd):/app -w=/app -m=1000m golang:1.22 go test -v ./... -expected=1048576000 -cgroup-version 2 - name: Run tests in Go container (4321m) run: | docker run --rm -v=$(pwd):/app -w=/app -m=4321m golang:1.22 go test -v ./... -expected=4530896896 -cgroup-version 2 - name: Run tests in Go container (system memory limit) run: | docker run --rm -v=$(pwd):/app -w=/app golang:1.22 go test -v ./... -expected-system=$(($(awk '/MemTotal/ {print $2}' /proc/meminfo) * 1024)) -cgroup-version 2 golang-github-kimmachinegun-automemlimit-0.7.1/.gitignore000066400000000000000000000004151475655524500236010ustar00rootroot00000000000000# Binaries for programs and plugins *.exe *.exe~ *.dll *.so *.dylib # Test binary, built with `go test -c` *.test # Output of the go coverage tool, specifically when used with LiteIDE *.out # Dependency directories (remove the comment below to include it) # vendor/ golang-github-kimmachinegun-automemlimit-0.7.1/LICENSE000066400000000000000000000020511475655524500226140ustar00rootroot00000000000000MIT 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. golang-github-kimmachinegun-automemlimit-0.7.1/README.md000066400000000000000000000046161475655524500230770ustar00rootroot00000000000000# automemlimit [![Go Reference](https://pkg.go.dev/badge/github.com/KimMachineGun/automemlimit.svg)](https://pkg.go.dev/github.com/KimMachineGun/automemlimit) [![Go Report Card](https://goreportcard.com/badge/github.com/KimMachineGun/automemlimit)](https://goreportcard.com/report/github.com/KimMachineGun/automemlimit) [![Test](https://github.com/KimMachineGun/automemlimit/actions/workflows/test.yml/badge.svg?branch=main)](https://github.com/KimMachineGun/automemlimit/actions/workflows/test.yml) Automatically set `GOMEMLIMIT` to match Linux [cgroups(7)](https://man7.org/linux/man-pages/man7/cgroups.7.html) memory limit. See more details about `GOMEMLIMIT` [here](https://tip.golang.org/doc/gc-guide#Memory_limit). ## Notice Version `v0.5.0` introduces a fallback to system memory limits as an experimental feature when cgroup limits are unavailable. Activate this by setting `AUTOMEMLIMIT_EXPERIMENT=system`. You can also use system memory limits via `memlimit.FromSystem` provider directly. This feature is under evaluation and might become a default or be removed based on user feedback. If you have any feedback about this feature, please open an issue. ## Installation ```shell go get github.com/KimMachineGun/automemlimit@latest ``` ## Usage ```go package main // By default, it sets `GOMEMLIMIT` to 90% of cgroup's memory limit. // This is equivalent to `memlimit.SetGoMemLimitWithOpts(memlimit.WithLogger(slog.Default()))` // To disable logging, use `memlimit.SetGoMemLimitWithOpts` directly. import _ "github.com/KimMachineGun/automemlimit" ``` or ```go package main import "github.com/KimMachineGun/automemlimit/memlimit" func init() { memlimit.SetGoMemLimitWithOpts( memlimit.WithRatio(0.9), memlimit.WithProvider(memlimit.FromCgroup), memlimit.WithLogger(slog.Default()), memlimit.WithRefreshInterval(1*time.Minute), ) memlimit.SetGoMemLimitWithOpts( memlimit.WithRatio(0.9), memlimit.WithProvider( memlimit.ApplyFallback( memlimit.FromCgroup, memlimit.FromSystem, ), ), memlimit.WithRefreshInterval(1*time.Minute), ) memlimit.SetGoMemLimit(0.9) memlimit.SetGoMemLimitWithProvider(memlimit.Limit(1024*1024), 0.9) memlimit.SetGoMemLimitWithProvider(memlimit.FromCgroup, 0.9) memlimit.SetGoMemLimitWithProvider(memlimit.FromCgroupV1, 0.9) memlimit.SetGoMemLimitWithProvider(memlimit.FromCgroupHybrid, 0.9) memlimit.SetGoMemLimitWithProvider(memlimit.FromCgroupV2, 0.9) } ``` golang-github-kimmachinegun-automemlimit-0.7.1/automemlimit.go000066400000000000000000000002741475655524500246510ustar00rootroot00000000000000package automemlimit import ( "log/slog" "github.com/KimMachineGun/automemlimit/memlimit" ) func init() { memlimit.SetGoMemLimitWithOpts( memlimit.WithLogger(slog.Default()), ) } golang-github-kimmachinegun-automemlimit-0.7.1/examples/000077500000000000000000000000001475655524500234275ustar00rootroot00000000000000golang-github-kimmachinegun-automemlimit-0.7.1/examples/dynamic/000077500000000000000000000000001475655524500250535ustar00rootroot00000000000000golang-github-kimmachinegun-automemlimit-0.7.1/examples/dynamic/go.mod000066400000000000000000000004351475655524500261630ustar00rootroot00000000000000module github.com/KimMachineGun/automemlimit/examples/dynamic go 1.22.0 toolchain go1.23.3 require github.com/KimMachineGun/automemlimit v0.0.0 require github.com/pbnjay/memory v0.0.0-20210728143218-7b4eea64cf58 // indirect replace github.com/KimMachineGun/automemlimit => ../../ golang-github-kimmachinegun-automemlimit-0.7.1/examples/dynamic/go.sum000066400000000000000000000003371475655524500262110ustar00rootroot00000000000000github.com/pbnjay/memory v0.0.0-20210728143218-7b4eea64cf58 h1:onHthvaw9LFnH4t2DcNVpwGmV9E1BkGknEliJkfwQj0= github.com/pbnjay/memory v0.0.0-20210728143218-7b4eea64cf58/go.mod h1:DXv8WO4yhMYhSNPKjeNKa5WY9YCIEBRbNzFFPJbWO6Y= golang-github-kimmachinegun-automemlimit-0.7.1/examples/dynamic/limit.txt000066400000000000000000000000131475655524500267240ustar00rootroot000000000000004294967296 golang-github-kimmachinegun-automemlimit-0.7.1/examples/dynamic/main.go000066400000000000000000000017021475655524500263260ustar00rootroot00000000000000package main import ( "bytes" "errors" "log/slog" "os" "os/signal" "strconv" "time" "github.com/KimMachineGun/automemlimit/memlimit" ) func init() { slog.SetDefault(slog.New(slog.NewJSONHandler(os.Stderr, nil))) memlimit.SetGoMemLimitWithOpts( memlimit.WithProvider( FileProvider("limit.txt"), ), memlimit.WithRefreshInterval(5*time.Second), memlimit.WithLogger(slog.Default()), ) } func main() { c := make(chan os.Signal, 1) signal.Notify(c, os.Interrupt) s := <-c slog.Info("signal captured", slog.Any("signal", s)) } func FileProvider(path string) memlimit.Provider { return func() (uint64, error) { b, err := os.ReadFile(path) if err != nil { if errors.Is(err, os.ErrNotExist) { return memlimit.ApplyFallback(memlimit.FromCgroup, memlimit.FromSystem)() } return 0, err } b = bytes.TrimSpace(b) if len(b) == 0 { return 0, memlimit.ErrNoLimit } return strconv.ParseUint(string(b), 10, 64) } } golang-github-kimmachinegun-automemlimit-0.7.1/examples/gosigar/000077500000000000000000000000001475655524500250625ustar00rootroot00000000000000golang-github-kimmachinegun-automemlimit-0.7.1/examples/gosigar/go.mod000066400000000000000000000011131475655524500261640ustar00rootroot00000000000000module github.com/KimMachineGun/automemlimit/examples/gosigar go 1.22.0 toolchain go1.23.3 require ( github.com/KimMachineGun/automemlimit v0.0.0 github.com/cloudfoundry/gosigar v1.3.30 ) require ( github.com/google/go-cmp v0.6.0 // indirect github.com/pbnjay/memory v0.0.0-20210728143218-7b4eea64cf58 // indirect github.com/pkg/errors v0.9.1 // indirect github.com/stretchr/testify v1.8.4 // indirect golang.org/x/net v0.31.0 // indirect golang.org/x/sys v0.27.0 // indirect golang.org/x/tools v0.27.0 // indirect ) replace github.com/KimMachineGun/automemlimit => ../../ golang-github-kimmachinegun-automemlimit-0.7.1/examples/gosigar/go.sum000066400000000000000000000056451475655524500262270ustar00rootroot00000000000000github.com/cloudfoundry/gosigar v1.3.30 h1:ZQPPt8RY72T8V+OZqPAi1qzkqH6UPhrAY8lfmDklNuI= github.com/cloudfoundry/gosigar v1.3.30/go.mod h1:v1aji1eOWmI6/v9T9Gd9ef1a2FEi9m9/25UnfHO0org= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/go-logr/logr v1.2.4 h1:g01GSCwiDw2xSZfjJ2/T9M+S6pFdcNtFYsp+Y43HYDQ= github.com/go-logr/logr v1.2.4/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A= github.com/go-task/slim-sprig v0.0.0-20230315185526-52ccab3ef572 h1:tfuBGBXKqDEevZMzYi5KSi8KkcZtzBcTgAUUtapy0OI= github.com/go-task/slim-sprig v0.0.0-20230315185526-52ccab3ef572/go.mod h1:9Pwr4B2jHnOSGXyyzV8ROjYa2ojvAY6HCGYYfMoC3Ls= github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI= github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= github.com/google/pprof v0.0.0-20230926050212-f7f687d19a98 h1:pUa4ghanp6q4IJHwE9RwLgmVFfReJN+KbQ8ExNEUUoQ= github.com/google/pprof v0.0.0-20230926050212-f7f687d19a98/go.mod h1:czg5+yv1E0ZGTi6S6vVK1mke0fV+FaUhNGcd6VRS9Ik= github.com/onsi/ginkgo/v2 v2.12.1 h1:uHNEO1RP2SpuZApSkel9nEh1/Mu+hmQe7Q+Pepg5OYA= github.com/onsi/ginkgo/v2 v2.12.1/go.mod h1:TE309ZR8s5FsKKpuB1YAQYBzCaAfUgatB/xlT/ETL/o= github.com/onsi/gomega v1.27.10 h1:naR28SdDFlqrG6kScpT8VWpu1xWY5nJRCF3XaYyBjhI= github.com/onsi/gomega v1.27.10/go.mod h1:RsS8tutOdbdgzbPtzzATp12yT7kM5I5aElG3evPbQ0M= github.com/pbnjay/memory v0.0.0-20210728143218-7b4eea64cf58 h1:onHthvaw9LFnH4t2DcNVpwGmV9E1BkGknEliJkfwQj0= github.com/pbnjay/memory v0.0.0-20210728143218-7b4eea64cf58/go.mod h1:DXv8WO4yhMYhSNPKjeNKa5WY9YCIEBRbNzFFPJbWO6Y= github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk= github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo= golang.org/x/net v0.31.0 h1:68CPQngjLL0r2AlUKiSxtQFKvzRVbnzLwMUn5SzcLHo= golang.org/x/net v0.31.0/go.mod h1:P4fl1q7dY2hnZFxEk4pPSkDHF+QqjitcnDjUQyMM+pM= golang.org/x/sys v0.27.0 h1:wBqf8DvsY9Y/2P8gAfPDEYNuS30J4lPHJxXSb/nJZ+s= golang.org/x/sys v0.27.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= golang.org/x/text v0.20.0 h1:gK/Kv2otX8gz+wn7Rmb3vT96ZwuoxnQlY+HlJVj7Qug= golang.org/x/text v0.20.0/go.mod h1:D4IsuqiFMhST5bX19pQ9ikHC2GsaKyk/oF+pn3ducp4= golang.org/x/tools v0.27.0 h1:qEKojBykQkQ4EynWy4S8Weg69NumxKdn40Fce3uc/8o= golang.org/x/tools v0.27.0/go.mod h1:sUi0ZgbwW9ZPAq26Ekut+weQPR5eIM6GQLQ1Yjm1H0Q= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= golang-github-kimmachinegun-automemlimit-0.7.1/examples/gosigar/main.go000066400000000000000000000006061475655524500263370ustar00rootroot00000000000000package main import ( "github.com/KimMachineGun/automemlimit/memlimit" sigar "github.com/cloudfoundry/gosigar" ) func init() { memlimit.SetGoMemLimitWithOpts( memlimit.WithProvider( memlimit.ApplyFallback( memlimit.FromCgroup, FromGoSigar, ), ), ) } func main() {} func FromGoSigar() (uint64, error) { var mem sigar.Mem err := mem.Get() return mem.Total, err } golang-github-kimmachinegun-automemlimit-0.7.1/examples/logger/000077500000000000000000000000001475655524500247065ustar00rootroot00000000000000golang-github-kimmachinegun-automemlimit-0.7.1/examples/logger/go.mod000066400000000000000000000004341475655524500260150ustar00rootroot00000000000000module github.com/KimMachineGun/automemlimit/examples/logger go 1.22.0 toolchain go1.23.3 require github.com/KimMachineGun/automemlimit v0.0.0 require github.com/pbnjay/memory v0.0.0-20210728143218-7b4eea64cf58 // indirect replace github.com/KimMachineGun/automemlimit => ../../ golang-github-kimmachinegun-automemlimit-0.7.1/examples/logger/go.sum000066400000000000000000000003371475655524500260440ustar00rootroot00000000000000github.com/pbnjay/memory v0.0.0-20210728143218-7b4eea64cf58 h1:onHthvaw9LFnH4t2DcNVpwGmV9E1BkGknEliJkfwQj0= github.com/pbnjay/memory v0.0.0-20210728143218-7b4eea64cf58/go.mod h1:DXv8WO4yhMYhSNPKjeNKa5WY9YCIEBRbNzFFPJbWO6Y= golang-github-kimmachinegun-automemlimit-0.7.1/examples/logger/main.go000066400000000000000000000004521475655524500261620ustar00rootroot00000000000000package main import ( "log/slog" "os" "github.com/KimMachineGun/automemlimit/memlimit" ) func init() { memlimit.SetGoMemLimitWithOpts( memlimit.WithProvider( memlimit.Limit(1024*1024*1024), ), memlimit.WithLogger(slog.New(slog.NewJSONHandler(os.Stderr, nil))), ) } func main() {} golang-github-kimmachinegun-automemlimit-0.7.1/examples/system/000077500000000000000000000000001475655524500247535ustar00rootroot00000000000000golang-github-kimmachinegun-automemlimit-0.7.1/examples/system/go.mod000066400000000000000000000004341475655524500260620ustar00rootroot00000000000000module github.com/KimMachineGun/automemlimit/examples/system go 1.22.0 toolchain go1.23.3 require github.com/KimMachineGun/automemlimit v0.0.0 require github.com/pbnjay/memory v0.0.0-20210728143218-7b4eea64cf58 // indirect replace github.com/KimMachineGun/automemlimit => ../../ golang-github-kimmachinegun-automemlimit-0.7.1/examples/system/go.sum000066400000000000000000000003371475655524500261110ustar00rootroot00000000000000github.com/pbnjay/memory v0.0.0-20210728143218-7b4eea64cf58 h1:onHthvaw9LFnH4t2DcNVpwGmV9E1BkGknEliJkfwQj0= github.com/pbnjay/memory v0.0.0-20210728143218-7b4eea64cf58/go.mod h1:DXv8WO4yhMYhSNPKjeNKa5WY9YCIEBRbNzFFPJbWO6Y= golang-github-kimmachinegun-automemlimit-0.7.1/examples/system/main.go000066400000000000000000000004011475655524500262210ustar00rootroot00000000000000package main import ( "github.com/KimMachineGun/automemlimit/memlimit" ) func init() { memlimit.SetGoMemLimitWithOpts( memlimit.WithProvider( memlimit.ApplyFallback( memlimit.FromCgroup, memlimit.FromSystem, ), ), ) } func main() {} golang-github-kimmachinegun-automemlimit-0.7.1/go.mod000066400000000000000000000002211475655524500227120ustar00rootroot00000000000000module github.com/KimMachineGun/automemlimit go 1.22.0 toolchain go1.23.3 require github.com/pbnjay/memory v0.0.0-20210728143218-7b4eea64cf58 golang-github-kimmachinegun-automemlimit-0.7.1/go.sum000066400000000000000000000003371475655524500227470ustar00rootroot00000000000000github.com/pbnjay/memory v0.0.0-20210728143218-7b4eea64cf58 h1:onHthvaw9LFnH4t2DcNVpwGmV9E1BkGknEliJkfwQj0= github.com/pbnjay/memory v0.0.0-20210728143218-7b4eea64cf58/go.mod h1:DXv8WO4yhMYhSNPKjeNKa5WY9YCIEBRbNzFFPJbWO6Y= golang-github-kimmachinegun-automemlimit-0.7.1/memlimit/000077500000000000000000000000001475655524500234265ustar00rootroot00000000000000golang-github-kimmachinegun-automemlimit-0.7.1/memlimit/cgroups.go000066400000000000000000000300071475655524500254370ustar00rootroot00000000000000package 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 } golang-github-kimmachinegun-automemlimit-0.7.1/memlimit/cgroups_linux.go000066400000000000000000000020271475655524500266570ustar00rootroot00000000000000//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 }) } golang-github-kimmachinegun-automemlimit-0.7.1/memlimit/cgroups_linux_test.go000066400000000000000000000026741475655524500277260ustar00rootroot00000000000000//go:build linux // +build linux package memlimit import ( "testing" ) func TestFromCgroup(t *testing.T) { if expected == 0 { t.Skip() } limit, err := FromCgroup() if cgVersion == 0 && err != ErrNoCgroup { t.Fatalf("FromCgroup() error = %v, wantErr %v", err, ErrNoCgroup) } if err != nil { t.Fatalf("FromCgroup() error = %v, wantErr %v", err, nil) } if limit != expected { t.Fatalf("FromCgroup() got = %v, want %v", limit, expected) } } func TestFromCgroupHybrid(t *testing.T) { if expected == 0 { t.Skip() } limit, err := FromCgroupHybrid() if cgVersion == 0 && err != ErrNoCgroup { t.Fatalf("FromCgroupHybrid() error = %v, wantErr %v", err, ErrNoCgroup) } if err != nil { t.Fatalf("FromCgroupHybrid() error = %v, wantErr %v", err, nil) } if limit != expected { t.Fatalf("FromCgroupHybrid() got = %v, want %v", limit, expected) } } func TestFromCgroupV1(t *testing.T) { if expected == 0 || cgVersion != 1 { t.Skip() } limit, err := FromCgroupV1() if err != nil { t.Fatalf("FromCgroupV1() error = %v, wantErr %v", err, nil) } if limit != expected { t.Fatalf("FromCgroupV1() got = %v, want %v", limit, expected) } } func TestFromCgroupV2(t *testing.T) { if expected == 0 || cgVersion != 2 { t.Skip() } limit, err := FromCgroupV2() if err != nil { t.Fatalf("FromCgroupV2() error = %v, wantErr %v", err, nil) } if limit != expected { t.Fatalf("FromCgroupV2() got = %v, want %v", limit, expected) } } golang-github-kimmachinegun-automemlimit-0.7.1/memlimit/cgroups_test.go000066400000000000000000000146611475655524500265060ustar00rootroot00000000000000package memlimit import ( "reflect" "testing" ) func TestParseMountInfoLine(t *testing.T) { tests := []struct { name string input string want mountInfo wantErr string }{ { name: "valid line with optional field", input: "36 35 98:0 /mnt1 /mnt2 rw,noatime master:1 - ext3 /dev/root rw,errors=continue", want: mountInfo{ Root: "/mnt1", MountPoint: "/mnt2", FilesystemType: "ext3", SuperOptions: "rw,errors=continue", }, }, { name: "valid line without optional field", input: "731 771 0:59 /sysrq-trigger /proc/sysrq-trigger ro,nosuid,nodev,noexec,relatime - proc proc rw", want: mountInfo{ Root: "/sysrq-trigger", MountPoint: "/proc/sysrq-trigger", FilesystemType: "proc", SuperOptions: "rw", }, }, { name: "valid line with minimal fields (no optional fields)", input: "25 1 0:22 / /dev rw - devtmpfs udev rw", want: mountInfo{ Root: "/", MountPoint: "/dev", FilesystemType: "devtmpfs", SuperOptions: "rw", }, }, { name: "no separator", input: "36 35 98:0 /mnt1 /mnt2 rw,noatime master:1 ext3 /dev/root rw,errors=continue", wantErr: `invalid separator`, }, { name: "not enough fields on left side", input: "36 35 98:0 /mnt1 /mnt2 - ext3 /dev/root rw,errors=continue", wantErr: `not enough fields before separator: [36 35 98:0 /mnt1 /mnt2]`, }, { name: "not enough fields on right side", input: "36 35 98:0 /mnt1 /mnt2 rw,noatime master:1 - ext3", wantErr: `not enough fields after separator: [ext3]`, }, { name: "too many fields on right side", input: "36 35 98:0 /mnt1 /mnt2 rw,noatime master:1 - ext3 /dev/root rw extra", wantErr: `too many fields after separator: [ext3 /dev/root rw extra]`, }, { name: "empty line", input: "", wantErr: `empty line`, }, { name: "6 fields on left side (no optional field), should add empty optional field", input: "100 1 8:2 / /data rw - ext4 /dev/sda2 rw,relatime", want: mountInfo{ Root: "/", MountPoint: "/data", FilesystemType: "ext4", SuperOptions: "rw,relatime", }, }, { name: "multiple optional fields on left side (issue #26)", input: "465 34 253:0 / / rw,relatime shared:409 master:1 - xfs /dev/mapper/fedora-root rw,seclabel,attr2,inode64,logbufs=8,logbsize=32k,noquota", want: mountInfo{ Root: "/", MountPoint: "/", FilesystemType: "xfs", SuperOptions: "rw,seclabel,attr2,inode64,logbufs=8,logbsize=32k,noquota", }, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { got, err := parseMountInfoLine(tt.input) if tt.wantErr != "" { if err == nil { t.Fatalf("expected an error containing %q, got nil", tt.wantErr) } if err.Error() != tt.wantErr { t.Fatalf("expected error containing %q, got %q", tt.wantErr, err.Error()) } return } if err != nil { t.Fatalf("unexpected error: %v", err) } if !reflect.DeepEqual(got, tt.want) { t.Fatalf("expected %+v, got %+v", tt.want, got) } }) } } func TestParseCgroupHierarchyLine(t *testing.T) { tests := []struct { name string input string want cgroupHierarchy wantErr string }{ { name: "valid line with multiple controllers", input: "5:cpuacct,cpu,cpuset:/daemons", want: cgroupHierarchy{ HierarchyID: "5", ControllerList: "cpuacct,cpu,cpuset", CgroupPath: "/daemons", }, }, { name: "valid line with no controllers (cgroup v2)", input: "0::/system.slice/docker.service", want: cgroupHierarchy{ HierarchyID: "0", ControllerList: "", CgroupPath: "/system.slice/docker.service", }, }, { name: "invalid line - only two fields", input: "5:cpuacct,cpu,cpuset", wantErr: "not enough fields: [5 cpuacct,cpu,cpuset]", }, { name: "invalid line - too many fields", input: "5:cpuacct,cpu:cpuset:/daemons:extra", wantErr: "too many fields: [5 cpuacct,cpu cpuset /daemons extra]", }, { name: "empty line", input: "", wantErr: "empty line", }, { name: "line with empty controller list but valid fields", input: "2::/my_cgroup", want: cgroupHierarchy{ HierarchyID: "2", ControllerList: "", CgroupPath: "/my_cgroup", }, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { got, err := parseCgroupHierarchyLine(tt.input) if tt.wantErr != "" { if err == nil { t.Fatalf("expected an error containing %q, got nil", tt.wantErr) } if err.Error() != tt.wantErr { t.Fatalf("expected error containing %q, got %q", tt.wantErr, err.Error()) } return } if err != nil { t.Fatalf("unexpected error: %v", err) } if !reflect.DeepEqual(got, tt.want) { t.Fatalf("expected %+v, got %+v", tt.want, got) } }) } } func TestResolveCgroupPath(t *testing.T) { tests := []struct { name string mountpoint string root string cgroupRelPath string want string wantErr string }{ { name: "exact match with both root and cgroupRelPath as '/'", mountpoint: "/fake/mount", root: "/", cgroupRelPath: "/", want: "/fake/mount", }, { name: "exact match with a non-root path", mountpoint: "/fake/mount", root: "/container0", cgroupRelPath: "/container0", want: "/fake/mount", }, { name: "valid subpath under root", mountpoint: "/fake/mount", root: "/container0", cgroupRelPath: "/container0/group1", want: "/fake/mount/group1", }, { name: "invalid cgroup path outside root", mountpoint: "/fake/mount", root: "/container0", cgroupRelPath: "/other_container", wantErr: "invalid cgroup path: /other_container is not under root /container0", }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { got, err := resolveCgroupPath(tt.mountpoint, tt.root, tt.cgroupRelPath) if tt.wantErr != "" { if err == nil { t.Fatalf("expected an error containing %q, got nil", tt.wantErr) } if err.Error() != tt.wantErr { t.Fatalf("expected error containing %q, got %q", tt.wantErr, err.Error()) } return } if err != nil { t.Fatalf("unexpected error: %v", err) } if got != tt.want { t.Fatalf("expected path %q, got %q", tt.want, got) } }) } } golang-github-kimmachinegun-automemlimit-0.7.1/memlimit/cgroups_unsupported.go000066400000000000000000000005431475655524500301110ustar00rootroot00000000000000//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 } golang-github-kimmachinegun-automemlimit-0.7.1/memlimit/cgroups_unsupported_test.go000066400000000000000000000022041475655524500311440ustar00rootroot00000000000000//go:build !linux // +build !linux package memlimit import ( "testing" ) func TestFromCgroup(t *testing.T) { limit, err := FromCgroup() if err != ErrCgroupsNotSupported { t.Fatalf("FromCgroup() error = %v, wantErr %v", err, ErrCgroupsNotSupported) } if limit != 0 { t.Fatalf("FromCgroup() got = %v, want %v", limit, 0) } } func TestFromCgroupV1(t *testing.T) { limit, err := FromCgroupV1() if err != ErrCgroupsNotSupported { t.Fatalf("FromCgroupV1() error = %v, wantErr %v", err, ErrCgroupsNotSupported) } if limit != 0 { t.Fatalf("FromCgroupV1() got = %v, want %v", limit, 0) } } func TestFromCgroupHybrid(t *testing.T) { limit, err := FromCgroupHybrid() if err != ErrCgroupsNotSupported { t.Fatalf("FromCgroupHybrid() error = %v, wantErr %v", err, ErrCgroupsNotSupported) } if limit != 0 { t.Fatalf("FromCgroupHybrid() got = %v, want %v", limit, 0) } } func TestFromCgroupV2(t *testing.T) { limit, err := FromCgroupV2() if err != ErrCgroupsNotSupported { t.Fatalf("FromCgroupV2() error = %v, wantErr %v", err, ErrCgroupsNotSupported) } if limit != 0 { t.Fatalf("FromCgroupV2() got = %v, want %v", limit, 0) } } golang-github-kimmachinegun-automemlimit-0.7.1/memlimit/exp_system.go000066400000000000000000000003631475655524500261570ustar00rootroot00000000000000package 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 } golang-github-kimmachinegun-automemlimit-0.7.1/memlimit/experiment.go000066400000000000000000000024401475655524500261350ustar00rootroot00000000000000package 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 } golang-github-kimmachinegun-automemlimit-0.7.1/memlimit/experiment_test.go000066400000000000000000000023311475655524500271730ustar00rootroot00000000000000package memlimit import ( "fmt" "os" "reflect" "testing" ) func TestParseExperiments(t *testing.T) { tests := []struct { name string env string want Experiments wantErr error }{ { name: "empty", env: "", want: Experiments{}, }, { name: "unknown", env: "unknown", want: Experiments{}, wantErr: fmt.Errorf("unknown AUTOMEMLIMIT_EXPERIMENT unknown"), }, { name: "none", env: "none", want: Experiments{}, }, { name: "none - with other", env: "system,none", want: Experiments{}, }, { name: "system", env: "system", want: Experiments{ System: true, }, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { exp, ok := os.LookupEnv(envAUTOMEMLIMIT_EXPERIMENT) t.Cleanup(func() { if ok { os.Setenv(envAUTOMEMLIMIT_EXPERIMENT, exp) } else { os.Unsetenv(envAUTOMEMLIMIT_EXPERIMENT) } }) os.Setenv("AUTOMEMLIMIT_EXPERIMENT", tt.env) exps, err := parseExperiments() if !reflect.DeepEqual(exps, tt.want) { t.Errorf("experiments= %#v, want %#v", exps, tt.want) } if !reflect.DeepEqual(err, tt.wantErr) { t.Errorf("err = %#v, want %#v", err, tt.wantErr) } }) } } golang-github-kimmachinegun-automemlimit-0.7.1/memlimit/logger.go000066400000000000000000000005751475655524500252430ustar00rootroot00000000000000package 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 } golang-github-kimmachinegun-automemlimit-0.7.1/memlimit/memlimit.go000066400000000000000000000173631475655524500256040ustar00rootroot00000000000000package 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 } } golang-github-kimmachinegun-automemlimit-0.7.1/memlimit/memlimit_linux_test.go000066400000000000000000000122661475655524500300570ustar00rootroot00000000000000//go:build linux // +build linux package memlimit import ( "flag" "math" "os" "runtime/debug" "testing" ) var ( cgVersion uint64 expected uint64 expectedSystem uint64 ) func TestMain(m *testing.M) { flag.Uint64Var(&expected, "expected", 0, "Expected cgroup's memory limit") flag.Uint64Var(&expectedSystem, "expected-system", 0, "Expected system memory limit") flag.Uint64Var(&cgVersion, "cgroup-version", 0, "Cgroup version") flag.Parse() os.Exit(m.Run()) } func TestSetGoMemLimit(t *testing.T) { type args struct { ratio float64 } tests := []struct { name string args args want int64 wantErr error gomemlimit int64 skip bool }{ { name: "0.5", args: args{ ratio: 0.5, }, want: int64(float64(expected) * 0.5), wantErr: nil, gomemlimit: int64(float64(expected) * 0.5), skip: expected == 0 || cgVersion == 0, }, { name: "0.9", args: args{ ratio: 0.9, }, want: int64(float64(expected) * 0.9), wantErr: nil, gomemlimit: int64(float64(expected) * 0.9), skip: expected == 0 || cgVersion == 0, }, { name: "Unavailable", args: args{ ratio: 0.9, }, want: 0, wantErr: ErrCgroupsNotSupported, gomemlimit: math.MaxInt64, skip: cgVersion != 0, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { if tt.skip { t.Skip() } t.Cleanup(func() { debug.SetMemoryLimit(math.MaxInt64) }) got, err := SetGoMemLimit(tt.args.ratio) if err != tt.wantErr { t.Errorf("SetGoMemLimit() error = %v, wantErr %v", err, tt.wantErr) return } if got != tt.want { t.Errorf("SetGoMemLimit() got = %v, want %v", got, tt.want) } if debug.SetMemoryLimit(-1) != tt.gomemlimit { t.Errorf("debug.SetMemoryLimit(-1) got = %v, want %v", debug.SetMemoryLimit(-1), tt.gomemlimit) } }) } } func TestSetGoMemLimitWithProvider_WithCgroupProvider(t *testing.T) { type args struct { provider Provider ratio float64 } tests := []struct { name string args args want int64 wantErr error gomemlimit int64 skip bool }{ { name: "FromCgroup", args: args{ provider: FromCgroup, ratio: 0.9, }, want: int64(float64(expected) * 0.9), wantErr: nil, gomemlimit: int64(float64(expected) * 0.9), skip: expected == 0 || cgVersion == 0, }, { name: "FromCgroup_Unavailable", args: args{ provider: FromCgroup, ratio: 0.9, }, want: 0, wantErr: ErrNoCgroup, gomemlimit: math.MaxInt64, skip: expected == 0 || cgVersion != 0, }, { name: "FromCgroupV1", args: args{ provider: FromCgroupV1, ratio: 0.9, }, want: int64(float64(expected) * 0.9), wantErr: nil, gomemlimit: int64(float64(expected) * 0.9), skip: expected == 0 || cgVersion != 1, }, { name: "FromCgroupHybrid", args: args{ provider: FromCgroupHybrid, ratio: 0.9, }, want: int64(float64(expected) * 0.9), wantErr: nil, gomemlimit: int64(float64(expected) * 0.9), skip: expected == 0 || cgVersion != 1, }, { name: "FromCgroupV2", args: args{ provider: FromCgroupV2, ratio: 0.9, }, want: int64(float64(expected) * 0.9), wantErr: nil, gomemlimit: int64(float64(expected) * 0.9), skip: expected == 0 || cgVersion != 2, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { if tt.skip { t.Skip() } t.Cleanup(func() { debug.SetMemoryLimit(math.MaxInt64) }) got, err := SetGoMemLimitWithProvider(tt.args.provider, tt.args.ratio) if err != tt.wantErr { t.Errorf("SetGoMemLimitWithProvider() error = %v, wantErr %v", err, tt.wantErr) return } if got != tt.want { t.Errorf("SetGoMemLimitWithProvider() got = %v, want %v", got, tt.want) } if debug.SetMemoryLimit(-1) != tt.gomemlimit { t.Errorf("debug.SetMemoryLimit(-1) got = %v, want %v", debug.SetMemoryLimit(-1), tt.gomemlimit) } }) } } func TestSetGoMemLimitWithProvider_WithSystemProvider(t *testing.T) { type args struct { provider Provider ratio float64 } tests := []struct { name string args args want int64 wantErr error gomemlimit int64 skip bool }{ { name: "FromSystem", args: args{ provider: FromSystem, ratio: 0.9, }, want: int64(float64(expectedSystem) * 0.9), wantErr: nil, gomemlimit: int64(float64(expectedSystem) * 0.9), skip: expectedSystem == 0, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { if tt.skip { t.Skip() } t.Cleanup(func() { debug.SetMemoryLimit(math.MaxInt64) }) got, err := SetGoMemLimitWithProvider(tt.args.provider, tt.args.ratio) if err != tt.wantErr { t.Errorf("SetGoMemLimitWithProvider() error = %v, wantErr %v", err, tt.wantErr) return } if got != tt.want { t.Errorf("SetGoMemLimitWithProvider() got = %v, want %v", got, tt.want) } if debug.SetMemoryLimit(-1) != tt.gomemlimit { t.Errorf("debug.SetMemoryLimit(-1) got = %v, want %v", debug.SetMemoryLimit(-1), tt.gomemlimit) } }) } } golang-github-kimmachinegun-automemlimit-0.7.1/memlimit/memlimit_test.go000066400000000000000000000144751475655524500266440ustar00rootroot00000000000000package memlimit import ( "fmt" "math" "runtime/debug" "sync/atomic" "testing" "time" ) func TestLimit(t *testing.T) { type args struct { limit uint64 } tests := []struct { name string args args want uint64 wantErr error }{ { name: "0bytes", args: args{ limit: 0, }, want: 0, wantErr: nil, }, { name: "1kib", args: args{ limit: 1024, }, want: 1024, wantErr: nil, }, { name: "1mib", args: args{ limit: 1024 * 1024, }, want: 1024 * 1024, wantErr: nil, }, { name: "1gib", args: args{ limit: 1024 * 1024 * 1024, }, want: 1024 * 1024 * 1024, wantErr: nil, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { got, err := Limit(tt.args.limit)() if err != tt.wantErr { t.Errorf("Limit() error = %v, wantErr %v", err, tt.wantErr) return } if got != tt.want { t.Errorf("Limit() got = %v, want %v", got, tt.want) } }) } } func TestSetGoMemLimitWithProvider(t *testing.T) { type args struct { provider Provider ratio float64 } tests := []struct { name string args args want int64 wantErr error gomemlimit int64 }{ { name: "Limit_0.5", args: args{ provider: Limit(1024 * 1024 * 1024), ratio: 0.5, }, want: 536870912, wantErr: nil, gomemlimit: 536870912, }, { name: "Limit_0.9", args: args{ provider: Limit(1024 * 1024 * 1024), ratio: 0.9, }, want: 966367641, wantErr: nil, gomemlimit: 966367641, }, { name: "Limit_0.9_math.MaxUint64", args: args{ provider: Limit(math.MaxUint64), ratio: 0.9, }, want: math.MaxInt64, wantErr: nil, gomemlimit: math.MaxInt64, }, { name: "Limit_0.9_math.MaxUint64", args: args{ provider: Limit(math.MaxUint64), ratio: 0.9, }, want: math.MaxInt64, wantErr: nil, gomemlimit: math.MaxInt64, }, { name: "Limit_0.45_math.MaxUint64", args: args{ provider: Limit(math.MaxUint64), ratio: 0.45, }, want: 8301034833169298432, wantErr: nil, gomemlimit: 8301034833169298432, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { t.Cleanup(func() { debug.SetMemoryLimit(math.MaxInt64) }) got, err := SetGoMemLimitWithProvider(tt.args.provider, tt.args.ratio) if err != tt.wantErr { t.Errorf("SetGoMemLimitWithProvider() error = %v, wantErr %v", err, tt.wantErr) return } if got != tt.want { t.Errorf("SetGoMemLimitWithProvider() got = %v, want %v", got, tt.want) } if debug.SetMemoryLimit(-1) != tt.gomemlimit { t.Errorf("debug.SetMemoryLimit(-1) got = %v, want %v", debug.SetMemoryLimit(-1), tt.gomemlimit) } }) } } func TestSetGoMemLimitWithOpts(t *testing.T) { tests := []struct { name string opts []Option want int64 wantErr error gomemlimit int64 }{ { name: "unknown error", opts: []Option{ WithProvider(func() (uint64, error) { return 0, fmt.Errorf("unknown error") }), }, want: 0, wantErr: fmt.Errorf("failed to set GOMEMLIMIT: unknown error"), gomemlimit: math.MaxInt64, }, { name: "ErrNoLimit", opts: []Option{ WithProvider(func() (uint64, error) { return 0, ErrNoLimit }), }, want: 0, wantErr: nil, gomemlimit: math.MaxInt64, }, { name: "wrapped ErrNoLimit", opts: []Option{ WithProvider(func() (uint64, error) { return 0, fmt.Errorf("wrapped: %w", ErrNoLimit) }), }, want: 0, wantErr: nil, gomemlimit: math.MaxInt64, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { got, err := SetGoMemLimitWithOpts(tt.opts...) if tt.wantErr != nil && err.Error() != tt.wantErr.Error() { t.Errorf("SetGoMemLimitWithOpts() error = %v, wantErr %v", err, tt.wantErr) return } if got != tt.want { t.Errorf("SetGoMemLimitWithOpts() got = %v, want %v", got, tt.want) } if debug.SetMemoryLimit(-1) != tt.gomemlimit { t.Errorf("debug.SetMemoryLimit(-1) got = %v, want %v", debug.SetMemoryLimit(-1), tt.gomemlimit) } }) } } func TestSetGoMemLimitWithOpts_rollbackOnPanic(t *testing.T) { t.Cleanup(func() { debug.SetMemoryLimit(math.MaxInt64) }) limit := int64(987654321) _ = debug.SetMemoryLimit(987654321) _, err := SetGoMemLimitWithOpts( WithProvider(func() (uint64, error) { debug.SetMemoryLimit(123456789) panic("panic") }), WithRatio(1), ) if err == nil { t.Error("SetGoMemLimitWithOpts() error = nil, want panic") } curr := debug.SetMemoryLimit(-1) if curr != limit { t.Errorf("debug.SetMemoryLimit(-1) got = %v, want %v", curr, limit) } } func TestSetGoMemLimitWithOpts_WithRefreshInterval(t *testing.T) { t.Cleanup(func() { debug.SetMemoryLimit(math.MaxInt64) }) var limit atomic.Int64 output, err := SetGoMemLimitWithOpts( WithProvider(func() (uint64, error) { l := limit.Load() if l == 0 { return 0, ErrNoLimit } return uint64(l), nil }), WithRatio(1), WithRefreshInterval(10*time.Millisecond), ) if err != nil { t.Errorf("SetGoMemLimitWithOpts() error = %v", err) } else if output != limit.Load() { t.Errorf("SetGoMemLimitWithOpts() got = %v, want %v", output, limit.Load()) } // 1. no limit curr := debug.SetMemoryLimit(-1) if curr != math.MaxInt64 { t.Errorf("debug.SetMemoryLimit(-1) got = %v, want %v", curr, limit.Load()) } // 2. max limit limit.Add(math.MaxInt64) time.Sleep(100 * time.Millisecond) curr = debug.SetMemoryLimit(-1) if curr != math.MaxInt64 { t.Errorf("debug.SetMemoryLimit(-1) got = %v, want %v", curr, int64(math.MaxInt64)) } // 3. adjust limit limit.Add(-1024) time.Sleep(100 * time.Millisecond) curr = debug.SetMemoryLimit(-1) if curr != math.MaxInt64-1024 { t.Errorf("debug.SetMemoryLimit(-1) got = %v, want %v", curr, int64(math.MaxInt64)-1024) } // 4. no limit again limit.Store(0) time.Sleep(100 * time.Millisecond) curr = debug.SetMemoryLimit(-1) if curr != math.MaxInt64 { t.Errorf("debug.SetMemoryLimit(-1) got = %v, want %v", curr, int64(math.MaxInt64)) } // 5. new limit limit.Store(math.MaxInt32) time.Sleep(100 * time.Millisecond) curr = debug.SetMemoryLimit(-1) if curr != math.MaxInt32 { t.Errorf("debug.SetMemoryLimit(-1) got = %v, want %v", curr, math.MaxInt32) } } golang-github-kimmachinegun-automemlimit-0.7.1/memlimit/memlimit_unsupported_test.go000066400000000000000000000072241475655524500313060ustar00rootroot00000000000000//go:build !linux // +build !linux package memlimit import ( "errors" "flag" "math" "os" "runtime/debug" "testing" ) var expected uint64 func TestMain(m *testing.M) { flag.Uint64Var(&expected, "expected", 0, "Expected memory limit") flag.Parse() os.Exit(m.Run()) } func TestSetGoMemLimit(t *testing.T) { type args struct { ratio float64 } tests := []struct { name string args args want int64 wantErr error gomemlimit int64 }{ { name: "0.5", args: args{ ratio: 0.5, }, want: 0, wantErr: ErrCgroupsNotSupported, gomemlimit: math.MaxInt64, }, { name: "0.9", args: args{ ratio: 0.9, }, want: 0, wantErr: ErrCgroupsNotSupported, gomemlimit: math.MaxInt64, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { t.Cleanup(func() { debug.SetMemoryLimit(math.MaxInt64) }) got, err := SetGoMemLimit(tt.args.ratio) if !errors.Is(err, tt.wantErr) { t.Errorf("SetGoMemLimit() error = %v, wantErr %v", err, tt.wantErr) return } if got != tt.want { t.Errorf("SetGoMemLimit() got = %v, want %v", got, tt.want) } if debug.SetMemoryLimit(-1) != tt.gomemlimit { t.Errorf("debug.SetMemoryLimit(-1) got = %v, want %v", debug.SetMemoryLimit(-1), tt.gomemlimit) } }) } } func TestSetGoMemLimitWithProvider_WithCgroupProvider(t *testing.T) { type args struct { provider Provider ratio float64 } tests := []struct { name string args args want int64 wantErr error }{ { name: "FromCgroup", args: args{ provider: FromCgroup, ratio: 0.9, }, want: 0, wantErr: ErrCgroupsNotSupported, }, { name: "FromCgroupV1", args: args{ provider: FromCgroupV1, ratio: 0.9, }, want: 0, wantErr: ErrCgroupsNotSupported, }, { name: "FromCgroupHybrid", args: args{ provider: FromCgroupHybrid, ratio: 0.9, }, want: 0, wantErr: ErrCgroupsNotSupported, }, { name: "FromCgroupV2", args: args{ provider: FromCgroupV2, ratio: 0.9, }, want: 0, wantErr: ErrCgroupsNotSupported, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { got, err := SetGoMemLimitWithProvider(tt.args.provider, tt.args.ratio) if !errors.Is(err, tt.wantErr) { t.Errorf("SetGoMemLimitWithProvider() error = %v, wantErr %v", err, tt.wantErr) return } if got != tt.want { t.Errorf("SetGoMemLimitWithProvider() got = %v, want %v", got, tt.want) } }) } } func TestSetGoMemLimitWithProvider_WithSystemProvider(t *testing.T) { type args struct { provider Provider ratio float64 } tests := []struct { name string args args want int64 wantErr error gomemlimit int64 skip bool }{ { name: "FromSystem", args: args{ provider: FromSystem, ratio: 0.9, }, want: int64(float64(expected) * 0.9), wantErr: nil, gomemlimit: int64(float64(expected) * 0.9), skip: expected == 0, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { if tt.skip { t.Skip() } t.Cleanup(func() { debug.SetMemoryLimit(math.MaxInt64) }) got, err := SetGoMemLimitWithProvider(tt.args.provider, tt.args.ratio) if !errors.Is(err, tt.wantErr) { t.Errorf("SetGoMemLimitWithProvider() error = %v, wantErr %v", err, tt.wantErr) return } if got != tt.want { t.Errorf("SetGoMemLimitWithProvider() got = %v, want %v", got, tt.want) } if debug.SetMemoryLimit(-1) != tt.gomemlimit { t.Errorf("debug.SetMemoryLimit(-1) got = %v, want %v", debug.SetMemoryLimit(-1), tt.gomemlimit) } }) } } golang-github-kimmachinegun-automemlimit-0.7.1/memlimit/provider.go000066400000000000000000000020441475655524500256070ustar00rootroot00000000000000package 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 } }