From ff30de2464cd7e974b8b9c4ff16eb1ac43795ed9 Mon Sep 17 00:00:00 2001 From: mzack Date: Tue, 5 Mar 2024 01:00:20 +0100 Subject: [PATCH 1/9] memguardian - backpressure --- go.mod | 11 +++-- go.sum | 29 ++++++----- memguardian/memguardian.go | 99 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 121 insertions(+), 18 deletions(-) create mode 100644 memguardian/memguardian.go diff --git a/go.mod b/go.mod index 61d2bd35..721e9f52 100644 --- a/go.mod +++ b/go.mod @@ -22,14 +22,15 @@ require ( github.com/projectdiscovery/machineid v0.0.0-20240226150047-2e2c51e35983 github.com/remeh/sizedwaitgroup v1.0.0 github.com/saintfish/chardet v0.0.0-20230101081208-5e3ef4b5456d - github.com/shirou/gopsutil/v3 v3.23.7 + github.com/shirou/gopsutil v3.21.11+incompatible + github.com/shirou/gopsutil/v3 v3.24.2 github.com/stretchr/testify v1.8.4 github.com/tidwall/gjson v1.14.3 github.com/zmap/zcrypto v0.0.0-20230422215203-9a665e1e9968 go.uber.org/multierr v1.11.0 golang.org/x/exp v0.0.0-20221205204356-47842c84f3db golang.org/x/oauth2 v0.11.0 - golang.org/x/sys v0.16.0 + golang.org/x/sys v0.17.0 golang.org/x/text v0.14.0 gopkg.in/yaml.v3 v3.0.1 ) @@ -83,15 +84,15 @@ require ( github.com/tidwall/pretty v1.2.0 // indirect github.com/tidwall/rtred v0.1.2 // indirect github.com/tidwall/tinyqueue v0.1.1 // indirect - github.com/tklauser/go-sysconf v0.3.11 // indirect - github.com/tklauser/numcpus v0.6.0 // indirect + github.com/tklauser/go-sysconf v0.3.12 // indirect + github.com/tklauser/numcpus v0.6.1 // indirect github.com/ulikunitz/xz v0.5.11 // indirect github.com/ulule/deepcopier v0.0.0-20200430083143-45decc6639b6 // indirect github.com/xi2/xz v0.0.0-20171230120015-48954b6210f8 // indirect github.com/yl2chen/cidranger v1.0.2 // indirect github.com/yuin/goldmark v1.5.4 // indirect github.com/yuin/goldmark-emoji v1.0.1 // indirect - github.com/yusufpapurcu/wmi v1.2.3 // indirect + github.com/yusufpapurcu/wmi v1.2.4 // indirect go.etcd.io/bbolt v1.3.7 // indirect gopkg.in/djherbis/times.v1 v1.3.0 // indirect ) diff --git a/go.sum b/go.sum index a22215a6..b2e7015e 100644 --- a/go.sum +++ b/go.sum @@ -82,8 +82,9 @@ github.com/google/go-cmp v0.5.2/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/ github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= github.com/google/go-cmp v0.5.6/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= github.com/google/go-cmp v0.5.8/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= -github.com/google/go-cmp v0.5.9 h1:O2Tfq5qg4qc4AmwVlvv0oLiVAGB7enBSJ2x2DqQFi38= github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= +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/go-github/v30 v30.1.0 h1:VLDx+UolQICEOKu2m4uAoMti1SxuEBAl7RSEG16L+Oo= github.com/google/go-github/v30 v30.1.0/go.mod h1:n8jBpHl45a/rlBUtRJMOG4GhNADUQFEufcolZ95JfU8= github.com/google/go-github/v50 v50.1.0/go.mod h1:Ev4Tre8QoKiolvbpOSG3FIi4Mlon3S2Nt9W5JYqKiwA= @@ -212,8 +213,10 @@ github.com/rivo/uniseg v0.4.4 h1:8TfxU8dW6PdqD27gjM8MVNuicgxIjxpm4K7x4jp8sis= github.com/rivo/uniseg v0.4.4/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88= github.com/saintfish/chardet v0.0.0-20230101081208-5e3ef4b5456d h1:hrujxIzL1woJ7AwssoOcM/tq5JjjG2yYOc8odClEiXA= github.com/saintfish/chardet v0.0.0-20230101081208-5e3ef4b5456d/go.mod h1:uugorj2VCxiV1x+LzaIdVa9b4S4qGAcH6cbhh4qVxOU= -github.com/shirou/gopsutil/v3 v3.23.7 h1:C+fHO8hfIppoJ1WdsVm1RoI0RwXoNdfTK7yWXV0wVj4= -github.com/shirou/gopsutil/v3 v3.23.7/go.mod h1:c4gnmoRC0hQuaLqvxnx1//VXQ0Ms/X9UnJF8pddY5z4= +github.com/shirou/gopsutil v3.21.11+incompatible h1:+1+c1VGhc88SSonWP6foOcLhvnKlUeu/erjjvaPEYiI= +github.com/shirou/gopsutil v3.21.11+incompatible/go.mod h1:5b4v6he4MtMOwMlS0TUMTu2PcXUg8+E1lC7eC3UO/RA= +github.com/shirou/gopsutil/v3 v3.24.2 h1:kcR0erMbLg5/3LcInpw0X/rrPSqq4CDPyI6A6ZRC18Y= +github.com/shirou/gopsutil/v3 v3.24.2/go.mod h1:tSg/594BcA+8UdQU2XcW803GWYgdtauFFPgJCJKZlVk= github.com/shoenig/go-m1cpu v0.1.6 h1:nxdKQNcEB6vzgA2E2bvzKIYRuNj7XNJ4S/aRSwKzFtM= github.com/shoenig/go-m1cpu v0.1.6/go.mod h1:1JJMcUBvfNwpq05QDQVAnx3gUHr9IYF7GNg9SUEw2VQ= github.com/shoenig/test v0.6.4 h1:kVTaSd7WLz5WZ2IaoM0RSzRsUD+m8wRR+5qvntpn4LU= @@ -258,10 +261,10 @@ github.com/tidwall/rtred v0.1.2 h1:exmoQtOLvDoO8ud++6LwVsAMTu0KPzLTUrMln8u1yu8= github.com/tidwall/rtred v0.1.2/go.mod h1:hd69WNXQ5RP9vHd7dqekAz+RIdtfBogmglkZSRxCHFQ= github.com/tidwall/tinyqueue v0.1.1 h1:SpNEvEggbpyN5DIReaJ2/1ndroY8iyEGxPYxoSaymYE= github.com/tidwall/tinyqueue v0.1.1/go.mod h1:O/QNHwrnjqr6IHItYrzoHAKYhBkLI67Q096fQP5zMYw= -github.com/tklauser/go-sysconf v0.3.11 h1:89WgdJhk5SNwJfu+GKyYveZ4IaJ7xAkecBo+KdJV0CM= -github.com/tklauser/go-sysconf v0.3.11/go.mod h1:GqXfhXY3kiPa0nAXPDIQIWzJbMCB7AmcWpGR8lSZfqI= -github.com/tklauser/numcpus v0.6.0 h1:kebhY2Qt+3U6RNK7UqpYNA+tJ23IBEGKkB7JQBfDYms= -github.com/tklauser/numcpus v0.6.0/go.mod h1:FEZLMke0lhOUG6w2JadTzp0a+Nl8PF/GFkQ5UVIcaL4= +github.com/tklauser/go-sysconf v0.3.12 h1:0QaGUFOdQaIVdPgfITYzaTegZvdCjmYO52cSFAEVmqU= +github.com/tklauser/go-sysconf v0.3.12/go.mod h1:Ho14jnntGE1fpdOqQEEaiKRpvIavV0hSfmBq8nJbHYI= +github.com/tklauser/numcpus v0.6.1 h1:ng9scYS7az0Bk4OZLvrNXNSAO2Pxr1XXRAPyjhIx+Fk= +github.com/tklauser/numcpus v0.6.1/go.mod h1:1XfjsgE2zo8GVw7POkMbHENHzVg3GzmoZ9fESEdAacY= github.com/ulikunitz/xz v0.5.8/go.mod h1:nbz6k7qbPmH4IRqmfOplQw/tblSgqTqBwxkY0oWt/14= github.com/ulikunitz/xz v0.5.9/go.mod h1:nbz6k7qbPmH4IRqmfOplQw/tblSgqTqBwxkY0oWt/14= github.com/ulikunitz/xz v0.5.11 h1:kpFauv27b6ynzBNT/Xy+1k+fK4WswhN/6PN5WhFAGw8= @@ -282,8 +285,8 @@ github.com/yuin/goldmark v1.5.4 h1:2uY/xC0roWy8IBEGLgB1ywIoEJFGmRrX21YQcvGZzjU= github.com/yuin/goldmark v1.5.4/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY= github.com/yuin/goldmark-emoji v1.0.1 h1:ctuWEyzGBwiucEqxzwe0SOYDXPAucOrE9NQC18Wa1os= github.com/yuin/goldmark-emoji v1.0.1/go.mod h1:2w1E6FEWLcDQkoTE+7HU6QF1F6SLlNGjRIBbIZQFqkQ= -github.com/yusufpapurcu/wmi v1.2.3 h1:E1ctvB7uKFMOJw3fdOW32DwGE9I7t++CRUEMKvFoFiw= -github.com/yusufpapurcu/wmi v1.2.3/go.mod h1:SBZ9tNy3G9/m5Oi98Zks0QjeHVDvuK0qfxQmPyzfmi0= +github.com/yusufpapurcu/wmi v1.2.4 h1:zFUKzehAFReQwLys1b/iSMl+JQGSCSjtVqQn9bBrPo0= +github.com/yusufpapurcu/wmi v1.2.4/go.mod h1:SBZ9tNy3G9/m5Oi98Zks0QjeHVDvuK0qfxQmPyzfmi0= github.com/zmap/rc2 v0.0.0-20131011165748-24b9757f5521/go.mod h1:3YZ9o3WnatTIZhuOtot4IcUfzoKVjUHqu6WALIyI0nE= github.com/zmap/rc2 v0.0.0-20190804163417-abaa70531248 h1:Nzukz5fNOBIHOsnP+6I79kPx3QhLv8nBy2mfFhBRq30= github.com/zmap/rc2 v0.0.0-20190804163417-abaa70531248/go.mod h1:3YZ9o3WnatTIZhuOtot4IcUfzoKVjUHqu6WALIyI0nE= @@ -370,12 +373,12 @@ golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f/go.mod h1:oPkhp1MJrh7nUepCBc golang.org/x/sys v0.0.0-20220728004956-3c1f35247d10/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20220908164124-27713097b956/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= -golang.org/x/sys v0.2.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= -golang.org/x/sys v0.10.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= -golang.org/x/sys v0.16.0 h1:xWw16ngr6ZMtmxDyKyIgsE93KNKz5HKmMa3b8ALHidU= -golang.org/x/sys v0.16.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= +golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.11.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.17.0 h1:25cE3gD+tdBA7lp7QfhuV+rJiE9YXTcS3VG1SqssI/Y= +golang.org/x/sys v0.17.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= golang.org/x/term v0.0.0-20201117132131-f5c789dd3221/go.mod h1:Nr5EML6q2oocZ2LXRh80K7BxOlk5/8JxuGnuhpl+muw= golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo= golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8= diff --git a/memguardian/memguardian.go b/memguardian/memguardian.go new file mode 100644 index 00000000..b5da8c07 --- /dev/null +++ b/memguardian/memguardian.go @@ -0,0 +1,99 @@ +package memguardian + +import ( + "context" + "errors" + "sync/atomic" + "time" + + "github.com/shirou/gopsutil/mem" +) + +type MemGuardianOption func(*MemGuardian) error + +func WitInterval(d time.Duration) MemGuardianOption { + return func(mg *MemGuardian) error { + mg.t = time.NewTicker(d) + return nil + } +} + +func WithCallback(f func()) MemGuardianOption { + return func(mg *MemGuardian) error { + mg.f = f + return nil + } +} + +func WithRatioWarning(ratio float64) MemGuardianOption { + return func(mg *MemGuardian) error { + if ratio == 0 || ratio > 100 { + return errors.New("ratio must be between 1 and 100") + } + mg.ratio = ratio + return nil + } +} + +type MemGuardian struct { + t *time.Ticker + f func() + ctx context.Context + cancel context.CancelFunc + Warning atomic.Bool + ratio float64 +} + +func New(options ...MemGuardianOption) (*MemGuardian, error) { + mg := &MemGuardian{} + for _, option := range options { + if err := option(mg); err != nil { + return nil, err + } + } + + mg.ctx, mg.cancel = context.WithCancel(context.TODO()) + + return mg, nil +} + +func (mg *MemGuardian) Run(ctx context.Context) error { + for { + select { + case <-mg.ctx.Done(): + mg.Close() + return nil + case <-ctx.Done(): + mg.Close() + return nil + case <-mg.t.C: + usedRatio, err := UsedRamRatio() + if err != nil { + return err + } + + if usedRatio >= mg.ratio { + mg.Warning.Store(true) + if mg.f != nil { + mg.f() + } + } else { + mg.Warning.Store(false) + } + } + } +} + +func (mg *MemGuardian) Close() { + mg.cancel() + mg.t.Stop() +} + +func UsedRamRatio() (float64, error) { + vms, err := mem.VirtualMemory() + if err != nil { + return 0, err + } + + return vms.UsedPercent, nil +} From 21775425fb10ab79b3734243806ae37638a3b02d Mon Sep 17 00:00:00 2001 From: mzack Date: Wed, 6 Mar 2024 13:18:46 +0100 Subject: [PATCH 2/9] adding comments --- memguardian/doc.go | 5 +++++ memguardian/memguardian.go | 7 +++++++ 2 files changed, 12 insertions(+) create mode 100644 memguardian/doc.go diff --git a/memguardian/doc.go b/memguardian/doc.go new file mode 100644 index 00000000..ef5180c0 --- /dev/null +++ b/memguardian/doc.go @@ -0,0 +1,5 @@ +// memguardian is a package that provides a simple RAM memory control mechanism +// once activated it sets an internal atomic boolean when the RAM usage exceed in absolute +// terms the warning ratio, for passive indirect check or invoke an optional callback for +// reactive backpressure +package memguardian diff --git a/memguardian/memguardian.go b/memguardian/memguardian.go index b5da8c07..0e15dea2 100644 --- a/memguardian/memguardian.go +++ b/memguardian/memguardian.go @@ -11,6 +11,7 @@ import ( type MemGuardianOption func(*MemGuardian) error +// WithInterval defines the ticker interval of the memory monitor func WitInterval(d time.Duration) MemGuardianOption { return func(mg *MemGuardian) error { mg.t = time.NewTicker(d) @@ -18,6 +19,7 @@ func WitInterval(d time.Duration) MemGuardianOption { } } +// WithCallback defines an optional callback if the warning ration is exceeded func WithCallback(f func()) MemGuardianOption { return func(mg *MemGuardian) error { mg.f = f @@ -25,6 +27,7 @@ func WithCallback(f func()) MemGuardianOption { } } +// WithRatioWarning defines the threshold of the warning state (and optional callback invocation) func WithRatioWarning(ratio float64) MemGuardianOption { return func(mg *MemGuardian) error { if ratio == 0 || ratio > 100 { @@ -44,6 +47,7 @@ type MemGuardian struct { ratio float64 } +// New mem guadian instance with user defined options func New(options ...MemGuardianOption) (*MemGuardian, error) { mg := &MemGuardian{} for _, option := range options { @@ -57,6 +61,7 @@ func New(options ...MemGuardianOption) (*MemGuardian, error) { return mg, nil } +// Run the instance monitor (cancel using the Stop method or context parameter) func (mg *MemGuardian) Run(ctx context.Context) error { for { select { @@ -84,11 +89,13 @@ func (mg *MemGuardian) Run(ctx context.Context) error { } } +// Close and stops the instance func (mg *MemGuardian) Close() { mg.cancel() mg.t.Stop() } +// Calculate the system absolute ratio of used RAM vs total available (as of now doesn't consider swap) func UsedRamRatio() (float64, error) { vms, err := mem.VirtualMemory() if err != nil { From 4400cc28a42f6c91e803d63cdb5a5ce69b3f360f Mon Sep 17 00:00:00 2001 From: mzack Date: Thu, 7 Mar 2024 15:56:30 +0100 Subject: [PATCH 3/9] creating default memory guard --- memguardian/memguardian.go | 82 +++++++++++++++++++++++++++++++------- 1 file changed, 67 insertions(+), 15 deletions(-) diff --git a/memguardian/memguardian.go b/memguardian/memguardian.go index 0e15dea2..1773315e 100644 --- a/memguardian/memguardian.go +++ b/memguardian/memguardian.go @@ -6,9 +6,45 @@ import ( "sync/atomic" "time" + units "github.com/docker/go-units" + "github.com/projectdiscovery/utils/env" "github.com/shirou/gopsutil/mem" ) +var ( + defaultInterval time.Duration + defaultMaxUsedRamRatio float64 + + DefaultMemGuardian *MemGuardian +) + +const ( + MemGuardianEnabled = "MEMGUARDIAN" + MemGuardianMaxUsedRamRatioENV = "MEMGUARDIAN_MAX_RAM_RATIO" + MemGuardianMaxUsedMemoryENV = "MEMGUARDIAN_MAX_RAM" + MemGuardianIntervalENV = "MEMGUARDIAN_INTERVAL" +) + +func init() { + defaultInterval = env.GetEnvOrDefault(MemGuardianMaxUsedRamRatioENV, time.Duration(time.Second*30)) + defaultMaxUsedRamRatio = env.GetEnvOrDefault(MemGuardianMaxUsedRamRatioENV, float64(75)) + maxRam := env.GetEnvOrDefault(MemGuardianMaxUsedRamRatioENV, "") + + options := []MemGuardianOption{ + WitInterval(defaultInterval), + WithMaxRamRatioWarning(defaultMaxUsedRamRatio), + } + if maxRam != "" { + options = append(options, WithMaxRamAmountWarning(maxRam)) + } + + var err error + DefaultMemGuardian, err = New(options...) + if err != nil { + panic(err) + } +} + type MemGuardianOption func(*MemGuardian) error // WithInterval defines the ticker interval of the memory monitor @@ -27,8 +63,8 @@ func WithCallback(f func()) MemGuardianOption { } } -// WithRatioWarning defines the threshold of the warning state (and optional callback invocation) -func WithRatioWarning(ratio float64) MemGuardianOption { +// WithMaxRamRatioWarning defines the ratio (1-100) threshold of the warning state (and optional callback invocation) +func WithMaxRamRatioWarning(ratio float64) MemGuardianOption { return func(mg *MemGuardian) error { if ratio == 0 || ratio > 100 { return errors.New("ratio must be between 1 and 100") @@ -38,13 +74,26 @@ func WithRatioWarning(ratio float64) MemGuardianOption { } } +// WithMaxRamAmountWarning defines the max amount of used RAM in bytes threshold of the warning state (and optional callback invocation) +func WithMaxRamAmountWarning(maxRam string) MemGuardianOption { + return func(mg *MemGuardian) error { + size, err := units.FromHumanSize(maxRam) + if err != nil { + return err + } + mg.maxMemory = uint64(size) + return nil + } +} + type MemGuardian struct { - t *time.Ticker - f func() - ctx context.Context - cancel context.CancelFunc - Warning atomic.Bool - ratio float64 + t *time.Ticker + f func() + ctx context.Context + cancel context.CancelFunc + Warning atomic.Bool + ratio float64 + maxMemory uint64 } // New mem guadian instance with user defined options @@ -72,12 +121,14 @@ func (mg *MemGuardian) Run(ctx context.Context) error { mg.Close() return nil case <-mg.t.C: - usedRatio, err := UsedRamRatio() + usedRatio, used, err := UsedRam() if err != nil { return err } - if usedRatio >= mg.ratio { + isRatioOverThreshold := mg.ratio > 0 && usedRatio >= mg.ratio + isAmountOverThreshold := mg.maxMemory > 0 && used >= mg.maxMemory + if isRatioOverThreshold || isAmountOverThreshold { mg.Warning.Store(true) if mg.f != nil { mg.f() @@ -95,12 +146,13 @@ func (mg *MemGuardian) Close() { mg.t.Stop() } -// Calculate the system absolute ratio of used RAM vs total available (as of now doesn't consider swap) -func UsedRamRatio() (float64, error) { - vms, err := mem.VirtualMemory() +// Calculate the system absolute ratio of used RAM +func UsedRam() (ratio float64, used uint64, err error) { + var vms *mem.VirtualMemoryStat + vms, err = mem.VirtualMemory() if err != nil { - return 0, err + return 0, 0, err } - return vms.UsedPercent, nil + return vms.UsedPercent, vms.Used, nil } From a623b4351f803d7174c08f9bd71d0579e9b69ecb Mon Sep 17 00:00:00 2001 From: mzack Date: Thu, 7 Mar 2024 17:20:43 +0100 Subject: [PATCH 4/9] adding docs --- memguardian/README.MD | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) create mode 100644 memguardian/README.MD diff --git a/memguardian/README.MD b/memguardian/README.MD new file mode 100644 index 00000000..7e9ff9d8 --- /dev/null +++ b/memguardian/README.MD @@ -0,0 +1,23 @@ +## Mem Guardian Usage Guide + +### Environment Variables + +- `MEMGUARDIAN`: Enable or disable memguardian. Set to 1 to enable +- `MEMGUARDIAN_MAX_RAM_RATIO`: Maximum ram ratio from 1 to 100 +- `MEMGUARDIAN_MAX_RAM`: Maximum amount of RAM (in size units ex: 10gb) +- `MEMGUARDIAN_INTERVAL`: detection interval (with unit ex: 30s) + + + +## How to Use + +1. Set the environment variables as per your requirements. + +```bash +export MEMGUARDIAN=1 +export MEMGUARDIAN_MAX_RAM_RATIO=75 # default +export MEMGUARDIAN_MAX_RAM=6Gb # optional +export MEMGUARDIAN_INTERVAL=30s # default +``` + +2. Run your Go application. The profiler will start automatically if MEMGUARDIAN is set to 1. \ No newline at end of file From 85520edc7b75ef632c61818acacfd11fa0316b53 Mon Sep 17 00:00:00 2001 From: mzack Date: Sat, 9 Mar 2024 22:40:31 +0100 Subject: [PATCH 5/9] using syscall on linux --- go.mod | 2 +- memguardian/memguardian.go | 6 ++---- memguardian/memory.go | 37 ++++++++++++++++++++++++++++++++++++ memguardian/memory_linux.go | 26 +++++++++++++++++++++++++ memguardian/memory_others.go | 23 ++++++++++++++++++++++ 5 files changed, 89 insertions(+), 5 deletions(-) create mode 100644 memguardian/memory.go create mode 100644 memguardian/memory_linux.go create mode 100644 memguardian/memory_others.go diff --git a/go.mod b/go.mod index 721e9f52..c32a16d0 100644 --- a/go.mod +++ b/go.mod @@ -8,6 +8,7 @@ require ( github.com/asaskevich/govalidator v0.0.0-20230301143203-a9d515a09cc2 github.com/cespare/xxhash v1.1.0 github.com/charmbracelet/glamour v0.6.0 + github.com/docker/go-units v0.5.0 github.com/google/go-github/v30 v30.1.0 github.com/hdm/jarm-go v0.0.7 github.com/julienschmidt/httprouter v1.3.0 @@ -45,7 +46,6 @@ require ( github.com/cloudflare/circl v1.3.7 // indirect github.com/dimchansky/utfbom v1.1.1 // indirect github.com/dlclark/regexp2 v1.8.1 // indirect - github.com/docker/go-units v0.5.0 // indirect github.com/dsnet/compress v0.0.2-0.20210315054119-f66993602bf5 // indirect github.com/fatih/color v1.15.0 // indirect github.com/gaukas/godicttls v0.0.4 // indirect diff --git a/memguardian/memguardian.go b/memguardian/memguardian.go index 1773315e..c64c0e28 100644 --- a/memguardian/memguardian.go +++ b/memguardian/memguardian.go @@ -8,7 +8,6 @@ import ( units "github.com/docker/go-units" "github.com/projectdiscovery/utils/env" - "github.com/shirou/gopsutil/mem" ) var ( @@ -148,11 +147,10 @@ func (mg *MemGuardian) Close() { // Calculate the system absolute ratio of used RAM func UsedRam() (ratio float64, used uint64, err error) { - var vms *mem.VirtualMemoryStat - vms, err = mem.VirtualMemory() + si, err := GetSysInfo() if err != nil { return 0, 0, err } - return vms.UsedPercent, vms.Used, nil + return si.UsedPercent(), si.UsedRam(), nil } diff --git a/memguardian/memory.go b/memguardian/memory.go new file mode 100644 index 00000000..ee52d68c --- /dev/null +++ b/memguardian/memory.go @@ -0,0 +1,37 @@ +package memguardian + +type SysInfo struct { + Uptime int64 + totalRam uint64 + freeRam uint64 + SharedRam uint64 + BufferRam uint64 + TotalSwap uint64 + FreeSwap uint64 + Unit uint64 + usedPercent float64 +} + +func (si *SysInfo) TotalRam() uint64 { + return uint64(si.totalRam) * uint64(si.Unit) +} + +func (si *SysInfo) FreeRam() uint64 { + return uint64(si.freeRam) * uint64(si.Unit) +} + +func (si *SysInfo) UsedRam() uint64 { + return si.TotalRam() - si.FreeRam() +} + +func (si *SysInfo) UsedPercent() float64 { + if si.usedPercent > 0 { + return si.usedPercent + } + + return 100 * float64((si.TotalRam()-si.FreeRam())*si.Unit) / float64(si.TotalRam()) +} + +func GetSysInfo() (*SysInfo, error) { + return getSysInfo() +} diff --git a/memguardian/memory_linux.go b/memguardian/memory_linux.go new file mode 100644 index 00000000..834e59be --- /dev/null +++ b/memguardian/memory_linux.go @@ -0,0 +1,26 @@ +//go:build linux + +package memguardian + +import "syscall" + +func getSysInfo() (*SysInfo, error) { + var sysInfo syscall.Sysinfo_t + err := syscall.Sysinfo(&sysInfo) + if err != nil { + return nil, err + } + + si := &SysInfo{ + Uptime: sysInfo.Uptime, + totalRam: sysInfo.Totalram, + freeRam: sysInfo.Freeram, + SharedRam: sysInfo.Freeram, + BufferRam: sysInfo.Bufferram, + TotalSwap: sysInfo.Totalswap, + FreeSwap: sysInfo.Freeswap, + Unit: uint64(sysInfo.Unit), + } + + return si, nil +} diff --git a/memguardian/memory_others.go b/memguardian/memory_others.go new file mode 100644 index 00000000..5d596d14 --- /dev/null +++ b/memguardian/memory_others.go @@ -0,0 +1,23 @@ +//go:build !linux + +package memguardian + +import "github.com/shirou/gopsutil/mem" + +// TODO: replace with native syscall +func getSysInfo() (*SysInfo, error) { + vms, err := mem.VirtualMemory() + if err != nil { + return nil, err + } + si := &SysInfo{ + totalRam: vms.Total, + freeRam: vms.Free, + SharedRam: vms.Shared, + TotalSwap: vms.SwapTotal, + FreeSwap: vms.SwapFree, + UsedPercent: vms.UsedPercent, + } + + return si, nil +} From 129e51f92c194f2d53ed8a1aae2bc12187c48516 Mon Sep 17 00:00:00 2001 From: mzack Date: Sat, 9 Mar 2024 22:46:54 +0100 Subject: [PATCH 6/9] . --- memguardian/memory_others.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/memguardian/memory_others.go b/memguardian/memory_others.go index 5d596d14..1a9bdf6d 100644 --- a/memguardian/memory_others.go +++ b/memguardian/memory_others.go @@ -16,7 +16,7 @@ func getSysInfo() (*SysInfo, error) { SharedRam: vms.Shared, TotalSwap: vms.SwapTotal, FreeSwap: vms.SwapFree, - UsedPercent: vms.UsedPercent, + usedPercent: vms.UsedPercent, } return si, nil From 5de745ca783c4c8c66172f8c8db2a56d7e42c91e Mon Sep 17 00:00:00 2001 From: mzack Date: Mon, 11 Mar 2024 19:47:16 +0100 Subject: [PATCH 7/9] exporting more fields --- http/respChain.go | 45 ++++++++++++++++++++++++++++++------- memguardian/memguardian.go | 12 +++++----- sync/semaphore/semaphore.go | 33 +++++++++++++++++++++------ sync/sizedpool/sizedpool.go | 5 +++++ 4 files changed, 74 insertions(+), 21 deletions(-) diff --git a/http/respChain.go b/http/respChain.go index d6b078fc..6a4e009b 100644 --- a/http/respChain.go +++ b/http/respChain.go @@ -2,25 +2,54 @@ package httputil import ( "bytes" + "context" "fmt" "net/http" "sync" + + "github.com/projectdiscovery/utils/sync/sizedpool" +) + +var ( + // reasonably high default allowed allocs + DefaultBytesBufferAlloc = int64(10000) ) +func ChangePoolSize(x int64) error { + return bufPool.Vary(context.Background(), x) +} + +func GetPoolSize() int64 { + return bufPool.Size() +} + // use buffer pool for storing response body // and reuse it for each request -var bufPool = sync.Pool{ - New: func() any { - // The Pool's New function should generally only return pointer - // types, since a pointer can be put into the return interface - // value without an allocation: - return new(bytes.Buffer) - }, +var bufPool *sizedpool.SizedPool[*bytes.Buffer] + +func init() { + var p = &sync.Pool{ + New: func() any { + // The Pool's New function should generally only return pointer + // types, since a pointer can be put into the return interface + // value without an allocation: + return new(bytes.Buffer) + }, + } + var err error + bufPool, err = sizedpool.New[*bytes.Buffer]( + sizedpool.WithPool[*bytes.Buffer](p), + sizedpool.WithSize[*bytes.Buffer](int64(DefaultBytesBufferAlloc)), + ) + if err != nil { + panic(err) + } } // getBuffer returns a buffer from the pool func getBuffer() *bytes.Buffer { - return bufPool.Get().(*bytes.Buffer) + buff, _ := bufPool.Get(context.Background()) + return buff } // putBuffer returns a buffer to the pool diff --git a/memguardian/memguardian.go b/memguardian/memguardian.go index c64c0e28..fd62c7cd 100644 --- a/memguardian/memguardian.go +++ b/memguardian/memguardian.go @@ -11,8 +11,8 @@ import ( ) var ( - defaultInterval time.Duration - defaultMaxUsedRamRatio float64 + DefaultInterval time.Duration + DefaultMaxUsedRamRatio float64 DefaultMemGuardian *MemGuardian ) @@ -25,13 +25,13 @@ const ( ) func init() { - defaultInterval = env.GetEnvOrDefault(MemGuardianMaxUsedRamRatioENV, time.Duration(time.Second*30)) - defaultMaxUsedRamRatio = env.GetEnvOrDefault(MemGuardianMaxUsedRamRatioENV, float64(75)) + DefaultInterval = env.GetEnvOrDefault(MemGuardianMaxUsedRamRatioENV, time.Duration(time.Second*30)) + DefaultMaxUsedRamRatio = env.GetEnvOrDefault(MemGuardianMaxUsedRamRatioENV, float64(75)) maxRam := env.GetEnvOrDefault(MemGuardianMaxUsedRamRatioENV, "") options := []MemGuardianOption{ - WitInterval(defaultInterval), - WithMaxRamRatioWarning(defaultMaxUsedRamRatio), + WitInterval(DefaultInterval), + WithMaxRamRatioWarning(DefaultMaxUsedRamRatio), } if maxRam != "" { options = append(options, WithMaxRamAmountWarning(maxRam)) diff --git a/sync/semaphore/semaphore.go b/sync/semaphore/semaphore.go index ffad47af..f4a355ea 100644 --- a/sync/semaphore/semaphore.go +++ b/sync/semaphore/semaphore.go @@ -4,24 +4,27 @@ import ( "context" "errors" "math" + "sync/atomic" "golang.org/x/sync/semaphore" ) type Semaphore struct { sem *semaphore.Weighted - initialSize int64 - maxSize int64 + initialSize atomic.Int64 + maxSize atomic.Int64 + currentSize atomic.Int64 } func New(size int64) (*Semaphore, error) { maxSize := int64(math.MaxInt64) s := &Semaphore{ - initialSize: size, - maxSize: maxSize, - sem: semaphore.NewWeighted(maxSize), + sem: semaphore.NewWeighted(maxSize), } - err := s.sem.Acquire(context.Background(), s.maxSize-s.initialSize) + s.initialSize.Store(size) + s.maxSize.Store(maxSize) + s.currentSize.Store(size) + err := s.sem.Acquire(context.Background(), s.maxSize.Load()-s.initialSize.Load()) return s, err } @@ -39,10 +42,26 @@ func (s *Semaphore) Vary(ctx context.Context, x int64) error { switch { case x > 0: s.sem.Release(x) + s.currentSize.Add(x) return nil case x < 0: - return s.sem.Acquire(ctx, x) + err := s.sem.Acquire(ctx, x) + if err != nil { + return err + } + s.currentSize.Add(x) + return nil default: return errors.New("x is zero") } } + +// Current size of the semaphore +func (s *Semaphore) Size() int64 { + return s.currentSize.Load() +} + +// Nominal size of the sempahore +func (s *Semaphore) InitialSize() int64 { + return s.initialSize.Load() +} diff --git a/sync/sizedpool/sizedpool.go b/sync/sizedpool/sizedpool.go index bfa9a696..3f70a46e 100644 --- a/sync/sizedpool/sizedpool.go +++ b/sync/sizedpool/sizedpool.go @@ -68,3 +68,8 @@ func (sz *SizedPool[T]) Put(x T) { func (sz *SizedPool[T]) Vary(ctx context.Context, x int64) error { return sz.sem.Vary(ctx, x) } + +// Current size of the pool +func (sz *SizedPool[T]) Size() int64 { + return sz.sem.Size() +} From 24c5b4be8fe1bf4b8921f2bef4f13a22a70096ed Mon Sep 17 00:00:00 2001 From: mzack Date: Mon, 11 Mar 2024 21:46:56 +0100 Subject: [PATCH 8/9] u64 --- memguardian/memory_linux.go | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/memguardian/memory_linux.go b/memguardian/memory_linux.go index 834e59be..207c3be0 100644 --- a/memguardian/memory_linux.go +++ b/memguardian/memory_linux.go @@ -12,13 +12,13 @@ func getSysInfo() (*SysInfo, error) { } si := &SysInfo{ - Uptime: sysInfo.Uptime, - totalRam: sysInfo.Totalram, - freeRam: sysInfo.Freeram, - SharedRam: sysInfo.Freeram, - BufferRam: sysInfo.Bufferram, - TotalSwap: sysInfo.Totalswap, - FreeSwap: sysInfo.Freeswap, + Uptime: int64(sysInfo.Uptime), + totalRam: uint64(sysInfo.Totalram), + freeRam: uint64(sysInfo.Freeram), + SharedRam: uint64(sysInfo.Freeram), + BufferRam: uint64(sysInfo.Bufferram), + TotalSwap: uint64(sysInfo.Totalswap), + FreeSwap: uint64(sysInfo.Freeswap), Unit: uint64(sysInfo.Unit), } From e23469b6173d4483f44e8c79e32428dff713b919 Mon Sep 17 00:00:00 2001 From: mzack Date: Wed, 13 Mar 2024 19:34:36 +0100 Subject: [PATCH 9/9] go.sum --- go.sum | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/go.sum b/go.sum index 8c887115..7b90717e 100644 --- a/go.sum +++ b/go.sum @@ -86,9 +86,8 @@ github.com/google/go-cmp v0.5.2/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/ github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= github.com/google/go-cmp v0.5.6/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= github.com/google/go-cmp v0.5.8/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= +github.com/google/go-cmp v0.5.9 h1:O2Tfq5qg4qc4AmwVlvv0oLiVAGB7enBSJ2x2DqQFi38= github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= -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/go-github/v30 v30.1.0 h1:VLDx+UolQICEOKu2m4uAoMti1SxuEBAl7RSEG16L+Oo= github.com/google/go-github/v30 v30.1.0/go.mod h1:n8jBpHl45a/rlBUtRJMOG4GhNADUQFEufcolZ95JfU8= github.com/google/go-github/v50 v50.1.0/go.mod h1:Ev4Tre8QoKiolvbpOSG3FIi4Mlon3S2Nt9W5JYqKiwA= @@ -219,8 +218,8 @@ github.com/saintfish/chardet v0.0.0-20230101081208-5e3ef4b5456d h1:hrujxIzL1woJ7 github.com/saintfish/chardet v0.0.0-20230101081208-5e3ef4b5456d/go.mod h1:uugorj2VCxiV1x+LzaIdVa9b4S4qGAcH6cbhh4qVxOU= github.com/shirou/gopsutil v3.21.11+incompatible h1:+1+c1VGhc88SSonWP6foOcLhvnKlUeu/erjjvaPEYiI= github.com/shirou/gopsutil v3.21.11+incompatible/go.mod h1:5b4v6he4MtMOwMlS0TUMTu2PcXUg8+E1lC7eC3UO/RA= -github.com/shirou/gopsutil/v3 v3.24.2 h1:kcR0erMbLg5/3LcInpw0X/rrPSqq4CDPyI6A6ZRC18Y= -github.com/shirou/gopsutil/v3 v3.24.2/go.mod h1:tSg/594BcA+8UdQU2XcW803GWYgdtauFFPgJCJKZlVk= +github.com/shirou/gopsutil/v3 v3.23.7 h1:C+fHO8hfIppoJ1WdsVm1RoI0RwXoNdfTK7yWXV0wVj4= +github.com/shirou/gopsutil/v3 v3.23.7/go.mod h1:c4gnmoRC0hQuaLqvxnx1//VXQ0Ms/X9UnJF8pddY5z4= github.com/shoenig/go-m1cpu v0.1.6 h1:nxdKQNcEB6vzgA2E2bvzKIYRuNj7XNJ4S/aRSwKzFtM= github.com/shoenig/go-m1cpu v0.1.6/go.mod h1:1JJMcUBvfNwpq05QDQVAnx3gUHr9IYF7GNg9SUEw2VQ= github.com/shoenig/test v0.6.4 h1:kVTaSd7WLz5WZ2IaoM0RSzRsUD+m8wRR+5qvntpn4LU= @@ -266,8 +265,10 @@ github.com/tidwall/rtred v0.1.2 h1:exmoQtOLvDoO8ud++6LwVsAMTu0KPzLTUrMln8u1yu8= github.com/tidwall/rtred v0.1.2/go.mod h1:hd69WNXQ5RP9vHd7dqekAz+RIdtfBogmglkZSRxCHFQ= github.com/tidwall/tinyqueue v0.1.1 h1:SpNEvEggbpyN5DIReaJ2/1ndroY8iyEGxPYxoSaymYE= github.com/tidwall/tinyqueue v0.1.1/go.mod h1:O/QNHwrnjqr6IHItYrzoHAKYhBkLI67Q096fQP5zMYw= +github.com/tklauser/go-sysconf v0.3.11/go.mod h1:GqXfhXY3kiPa0nAXPDIQIWzJbMCB7AmcWpGR8lSZfqI= github.com/tklauser/go-sysconf v0.3.12 h1:0QaGUFOdQaIVdPgfITYzaTegZvdCjmYO52cSFAEVmqU= github.com/tklauser/go-sysconf v0.3.12/go.mod h1:Ho14jnntGE1fpdOqQEEaiKRpvIavV0hSfmBq8nJbHYI= +github.com/tklauser/numcpus v0.6.0/go.mod h1:FEZLMke0lhOUG6w2JadTzp0a+Nl8PF/GFkQ5UVIcaL4= github.com/tklauser/numcpus v0.6.1 h1:ng9scYS7az0Bk4OZLvrNXNSAO2Pxr1XXRAPyjhIx+Fk= github.com/tklauser/numcpus v0.6.1/go.mod h1:1XfjsgE2zo8GVw7POkMbHENHzVg3GzmoZ9fESEdAacY= github.com/ulikunitz/xz v0.5.8/go.mod h1:nbz6k7qbPmH4IRqmfOplQw/tblSgqTqBwxkY0oWt/14= @@ -288,6 +289,7 @@ github.com/yuin/goldmark v1.5.4 h1:2uY/xC0roWy8IBEGLgB1ywIoEJFGmRrX21YQcvGZzjU= github.com/yuin/goldmark v1.5.4/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY= github.com/yuin/goldmark-emoji v1.0.1 h1:ctuWEyzGBwiucEqxzwe0SOYDXPAucOrE9NQC18Wa1os= github.com/yuin/goldmark-emoji v1.0.1/go.mod h1:2w1E6FEWLcDQkoTE+7HU6QF1F6SLlNGjRIBbIZQFqkQ= +github.com/yusufpapurcu/wmi v1.2.3/go.mod h1:SBZ9tNy3G9/m5Oi98Zks0QjeHVDvuK0qfxQmPyzfmi0= github.com/yusufpapurcu/wmi v1.2.4 h1:zFUKzehAFReQwLys1b/iSMl+JQGSCSjtVqQn9bBrPo0= github.com/yusufpapurcu/wmi v1.2.4/go.mod h1:SBZ9tNy3G9/m5Oi98Zks0QjeHVDvuK0qfxQmPyzfmi0= github.com/zmap/rc2 v0.0.0-20131011165748-24b9757f5521/go.mod h1:3YZ9o3WnatTIZhuOtot4IcUfzoKVjUHqu6WALIyI0nE= @@ -376,9 +378,11 @@ golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f/go.mod h1:oPkhp1MJrh7nUepCBc golang.org/x/sys v0.0.0-20220728004956-3c1f35247d10/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20220908164124-27713097b956/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.2.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.10.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.11.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.17.0 h1:25cE3gD+tdBA7lp7QfhuV+rJiE9YXTcS3VG1SqssI/Y= golang.org/x/sys v0.17.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=