diff --git a/go.mod b/go.mod index 5728489..96865a0 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/eapache/channels v1.1.0 github.com/google/go-github/v30 v30.1.0 github.com/hdm/jarm-go v0.0.7 @@ -23,6 +24,7 @@ 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.21.11+incompatible github.com/shirou/gopsutil/v3 v3.23.7 github.com/stretchr/testify v1.9.0 github.com/tidwall/gjson v1.14.3 @@ -31,7 +33,7 @@ require ( golang.org/x/exp v0.0.0-20221205204356-47842c84f3db golang.org/x/oauth2 v0.11.0 golang.org/x/sync v0.6.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 ) @@ -46,7 +48,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/eapache/queue v1.1.0 // indirect github.com/fatih/color v1.15.0 // indirect @@ -86,14 +87,14 @@ 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/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 49f3b34..7b90717 100644 --- a/go.sum +++ b/go.sum @@ -216,6 +216,8 @@ 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.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.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= @@ -263,10 +265,12 @@ 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/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= 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= @@ -285,8 +289,9 @@ 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= @@ -376,9 +381,11 @@ golang.org/x/sys v0.0.0-20220908164124-27713097b956/go.mod h1:oPkhp1MJrh7nUepCBc 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.16.0 h1:xWw16ngr6ZMtmxDyKyIgsE93KNKz5HKmMa3b8ALHidU= -golang.org/x/sys v0.16.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= +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/http/respChain.go b/http/respChain.go index d6b078f..6a4e009 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/README.MD b/memguardian/README.MD new file mode 100644 index 0000000..7e9ff9d --- /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 diff --git a/memguardian/doc.go b/memguardian/doc.go new file mode 100644 index 0000000..ef5180c --- /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 new file mode 100644 index 0000000..fd62c7c --- /dev/null +++ b/memguardian/memguardian.go @@ -0,0 +1,156 @@ +package memguardian + +import ( + "context" + "errors" + "sync/atomic" + "time" + + units "github.com/docker/go-units" + "github.com/projectdiscovery/utils/env" +) + +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 +func WitInterval(d time.Duration) MemGuardianOption { + return func(mg *MemGuardian) error { + mg.t = time.NewTicker(d) + return nil + } +} + +// WithCallback defines an optional callback if the warning ration is exceeded +func WithCallback(f func()) MemGuardianOption { + return func(mg *MemGuardian) error { + mg.f = f + return nil + } +} + +// 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") + } + mg.ratio = ratio + return nil + } +} + +// 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 + maxMemory uint64 +} + +// New mem guadian instance with user defined options +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 +} + +// Run the instance monitor (cancel using the Stop method or context parameter) +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, used, err := UsedRam() + if err != nil { + return err + } + + 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() + } + } else { + mg.Warning.Store(false) + } + } + } +} + +// Close and stops the instance +func (mg *MemGuardian) Close() { + mg.cancel() + mg.t.Stop() +} + +// Calculate the system absolute ratio of used RAM +func UsedRam() (ratio float64, used uint64, err error) { + si, err := GetSysInfo() + if err != nil { + return 0, 0, err + } + + return si.UsedPercent(), si.UsedRam(), nil +} diff --git a/memguardian/memory.go b/memguardian/memory.go new file mode 100644 index 0000000..ee52d68 --- /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 0000000..207c3be --- /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: 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), + } + + return si, nil +} diff --git a/memguardian/memory_others.go b/memguardian/memory_others.go new file mode 100644 index 0000000..1a9bdf6 --- /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 +} diff --git a/sync/semaphore/semaphore.go b/sync/semaphore/semaphore.go index ffad47a..f4a355e 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 bfa9a69..3f70a46 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() +}