-
Notifications
You must be signed in to change notification settings - Fork 33
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #362 from projectdiscovery/feat-memguardian
memguardian - backpressure
- Loading branch information
Showing
11 changed files
with
356 additions
and
25 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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. |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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 | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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() | ||
| } |
Oops, something went wrong.