split active+passive guards

go.mod

@@ -91,7 +91,7 @@ require (
 	github.com/projectdiscovery/tlsx v1.1.6
 	github.com/projectdiscovery/uncover v1.0.7
 	github.com/projectdiscovery/useragent v0.0.40
-	github.com/projectdiscovery/utils v0.0.83-0.20240305000020-ff30de2464cd
+	github.com/projectdiscovery/utils v0.0.84-0.20240311184716-5de745ca783c
 	github.com/projectdiscovery/wappalyzergo v0.0.112
 	github.com/redis/go-redis/v9 v9.1.0
 	github.com/sashabaranov/go-openai v1.15.3
@@ -217,6 +217,7 @@ require (
 	github.com/yuin/goldmark-emoji v1.0.1 // indirect
 	github.com/zeebo/blake3 v0.2.3 // indirect
 	golang.org/x/arch v0.3.0 // indirect
+	golang.org/x/sync v0.6.0 // indirect
 	gopkg.in/djherbis/times.v1 v1.3.0 // indirect
 	mellium.im/sasl v0.3.1 // indirect
 )

go.sum

@@ -864,10 +864,8 @@ github.com/projectdiscovery/uncover v1.0.7 h1:ut+2lTuvmftmveqF5RTjMWAgyLj8ltPQC7
 github.com/projectdiscovery/uncover v1.0.7/go.mod h1:HFXgm1sRPuoN0D4oATljPIdmbo/EEh1wVuxQqo/dwFE=
 github.com/projectdiscovery/useragent v0.0.40 h1:1LUhReSGPkhqsM5n40OOC9dIoNqMGs1dyGFJcOmg2Fo=
 github.com/projectdiscovery/useragent v0.0.40/go.mod h1:EvK1x3s948Gtqb/XOahXcauyejCL/rSgy5d1IAvsKT4=
-github.com/projectdiscovery/utils v0.0.83-0.20240305000020-ff30de2464cd h1:7AvCjcfZFkYy1Cg7aoA2iqVCuG/n+OHBFWB2ELGNAiI=
-github.com/projectdiscovery/utils v0.0.83-0.20240305000020-ff30de2464cd/go.mod h1:67zb3eUa96XlTpWDqnJhe7xVoDbyAwHb7ChkdooiQxQ=
-github.com/projectdiscovery/utils v0.0.83-0.20240307145630-4400cc28a42f h1:rryhFp/Ie6HjHAp8X0Ms0Z200f/COFjnH6+hAbF/2CQ=
-github.com/projectdiscovery/utils v0.0.83-0.20240307145630-4400cc28a42f/go.mod h1:67zb3eUa96XlTpWDqnJhe7xVoDbyAwHb7ChkdooiQxQ=
+github.com/projectdiscovery/utils v0.0.84-0.20240311184716-5de745ca783c h1:sRP8/40BeIkXpC5J3spdegkiLSDoyXHuP/OFYFYZsAc=
+github.com/projectdiscovery/utils v0.0.84-0.20240311184716-5de745ca783c/go.mod h1:paL6KzAPFZyaxgSWcd/0LHFGwLxFYoCghIAnaIMa+2k=
 github.com/projectdiscovery/wappalyzergo v0.0.112 h1:QPpp5jmj1lqLd5mFdFKQ9VvcYhQNqyU9Mr+IB0US2zA=
 github.com/projectdiscovery/wappalyzergo v0.0.112/go.mod h1:hc/o+fgM8KtdpFesjfBTmHTwsR+yBd+4kYZW/DGy/x8=
 github.com/projectdiscovery/yamldoc-go v1.0.4 h1:eZoESapnMw6WAHiVgRwNqvbJEfNHEH148uthhFbG5jE=

pkg/protocols/common/protocolstate/memguardian.go

@@ -1,23 +1,57 @@
 package protocolstate
 
 import (
+	"log"
+	"sync"
+	"time"
+
+	"github.com/projectdiscovery/gologger"
 	"github.com/projectdiscovery/utils/env"
+	httputil "github.com/projectdiscovery/utils/http"
 	"github.com/projectdiscovery/utils/memguardian"
 )
 
 var (
-	MaxThreadsOnLowMemory = env.GetEnvOrDefault("MEMGUARDIAN_THREADS", 0)
+	MaxThreadsOnLowMemory          = env.GetEnvOrDefault("MEMGUARDIAN_THREADS", 0)
+	MaxBytesBufferAllocOnLowMemory = env.GetEnvOrDefault("MEMGUARDIAN_ALLOC", 0)
+	memTimer                       *time.Ticker
 )
 
+func StartActiveMemGuardian() {
+	if memguardian.DefaultMemGuardian == nil {
+		return
+	}
+
+	memTimer := time.NewTicker(memguardian.DefaultInterval)
+	go func() {
+		for range memTimer.C {
+			log.Println(IsLowOnMemory())
+			if IsLowOnMemory() {
+				GlobalGuardBytesBufferAlloc()
+			} else {
+				GlobalRestoreBytesBufferAlloc()
+			}
+		}
+	}()
+}
+
+func StopActiveMemGuardian() {
+	if memguardian.DefaultMemGuardian == nil {
+		return
+	}
+
+	memTimer.Stop()
+}
+
 func IsLowOnMemory() bool {
 	if memguardian.DefaultMemGuardian != nil && memguardian.DefaultMemGuardian.Warning.Load() {
 		return true
 	}
-
 	return false
 }
 
-func GuardThreads(current int) int {
+// GuardThreads on caller
+func GuardThreadsOrDefault(current int) int {
 	if MaxThreadsOnLowMemory > 0 {
 		return MaxThreadsOnLowMemory
 	}
@@ -29,3 +63,38 @@ func GuardThreads(current int) int {
 
 	return 1
 }
+
+var muGlobalChange sync.Mutex
+
+// Global setting
+func GlobalGuardBytesBufferAlloc() error {
+	if muGlobalChange.TryLock() {
+		return nil
+
+	}
+	defer muGlobalChange.Unlock()
+
+	// if current capacity was not reduced decrease it
+	if MaxBytesBufferAllocOnLowMemory > 0 && httputil.DefaultBytesBufferAlloc == httputil.GetPoolSize() {
+		gologger.Debug().Msgf("reducing bytes.buffer pool size to: %d", MaxBytesBufferAllocOnLowMemory)
+		delta := httputil.GetPoolSize() - int64(MaxBytesBufferAllocOnLowMemory)
+		return httputil.ChangePoolSize(-delta)
+	}
+
+	return nil
+}
+
+// Global setting
+func GlobalRestoreBytesBufferAlloc() {
+	if muGlobalChange.TryLock() {
+		return
+
+	}
+	defer muGlobalChange.Unlock()
+
+	if httputil.DefaultBytesBufferAlloc != httputil.GetPoolSize() {
+		delta := httputil.DefaultBytesBufferAlloc - httputil.GetPoolSize()
+		gologger.Debug().Msgf("restoring bytes.buffer pool size to: %d", httputil.DefaultBytesBufferAlloc)
+		httputil.ChangePoolSize(delta)
+	}
+}

pkg/protocols/common/protocolstate/state.go

@@ -145,6 +145,8 @@ func Init(options *types.Options) error {
 		return Dialer.Dial(ctx, "tcp", addr)
 	})
 
+	StartActiveMemGuardian()
+
 	return nil
 }
 
@@ -205,4 +207,5 @@ func Close() {
 	if Dialer != nil {
 		Dialer.Close()
 	}
+	StopActiveMemGuardian()
 }

pkg/protocols/http/http.go

@@ -3,7 +3,6 @@ package http
 import (
 	"bytes"
 	"fmt"
-	"log"
 	"strings"
 
 	json "github.com/json-iterator/go"
@@ -392,11 +391,10 @@ func (request *Request) Compile(options *protocols.ExecutorOptions) error {
 	if len(request.Payloads) > 0 {
 		// specifically for http requests high concurrency and and threads will lead to memory exausthion, hence reduce the maximum parallelism
 		if protocolstate.IsLowOnMemory() {
-			request.Threads = protocolstate.GuardThreads(request.Threads)
+			request.Threads = protocolstate.GuardThreadsOrDefault(request.Threads)
 		}
 		// if we have payloads, adjust threads if none specified
 		request.Threads = options.GetThreadsForNPayloadRequests(request.Requests(), request.Threads)
-		log.Println("request.Threads:", request.Threads)
 	}
 
 	return nil

pkg/protocols/http/request.go

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