small change in calculation

This commit is contained in:
mzack 2024-03-07 17:28:24 +01:00
parent 89858a2ec8
commit 6f6c1d3679
3 changed files with 23 additions and 3 deletions

View File

@ -1,6 +1,13 @@
package protocolstate package protocolstate
import "github.com/projectdiscovery/utils/memguardian" import (
"github.com/projectdiscovery/utils/env"
"github.com/projectdiscovery/utils/memguardian"
)
var (
MaxThreadsOnLowMemory = env.GetEnvOrDefault("MEMGUARDIAN_THREADS", 0)
)
func IsLowOnMemory() bool { func IsLowOnMemory() bool {
if memguardian.DefaultMemGuardian != nil && memguardian.DefaultMemGuardian.Warning.Load() { if memguardian.DefaultMemGuardian != nil && memguardian.DefaultMemGuardian.Warning.Load() {
@ -9,3 +16,16 @@ func IsLowOnMemory() bool {
return false return false
} }
func GuardThreads(current int) int {
if MaxThreadsOnLowMemory > 0 {
return MaxThreadsOnLowMemory
}
fraction := int(current / 5)
if fraction > 0 {
return fraction
}
return 1
}

View File

@ -392,7 +392,7 @@ func (request *Request) Compile(options *protocols.ExecutorOptions) error {
if len(request.Payloads) > 0 { if len(request.Payloads) > 0 {
// specifically for http requests high concurrency and and threads will lead to memory exausthion, hence reduce the maximum parallelism // specifically for http requests high concurrency and and threads will lead to memory exausthion, hence reduce the maximum parallelism
if protocolstate.IsLowOnMemory() { if protocolstate.IsLowOnMemory() {
request.Threads = 5 request.Threads = protocolstate.GuardThreads(request.Threads)
} }
// if we have payloads, adjust threads if none specified // if we have payloads, adjust threads if none specified
request.Threads = options.GetThreadsForNPayloadRequests(request.Requests(), request.Threads) request.Threads = options.GetThreadsForNPayloadRequests(request.Requests(), request.Threads)

View File

@ -169,7 +169,7 @@ func (request *Request) executeParallelHTTP(input *contextargs.Context, dynamicV
maxWorkers := request.Threads maxWorkers := request.Threads
if protocolstate.IsLowOnMemory() { if protocolstate.IsLowOnMemory() {
maxWorkers = 5 maxWorkers = protocolstate.GuardThreads(request.Threads)
} }
// Stop-at-first-match logic while executing requests // Stop-at-first-match logic while executing requests