fix(protocolstate): use cgroup-aware memory guardian

go.mod

@@ -102,7 +102,7 @@ require (
 	github.com/projectdiscovery/tlsx v1.4.1
 	github.com/projectdiscovery/uncover v1.2.1
 	github.com/projectdiscovery/useragent v0.0.109
-	github.com/projectdiscovery/utils v0.11.9
+	github.com/projectdiscovery/utils v0.11.10-0.20261017183053-ac02d12acdf3
 	github.com/projectdiscovery/wappalyzergo v0.3.4
 	github.com/projectdiscovery/yamldoc-go v1.0.7
 	github.com/redis/go-redis/v9 v9.19.0

go.sum

@@ -1005,6 +1005,8 @@ github.com/projectdiscovery/useragent v0.0.109 h1:b86BPdZvmgbJousa04jAugDRY4Y+rA
 github.com/projectdiscovery/useragent v0.0.109/go.mod h1:sRyBqDAcD31RKYrO8fyXGowLd36fNwY1g0MNW9mp7zk=
 github.com/projectdiscovery/utils v0.11.9 h1:aKLje3zPcZuOGmopN51PasWxPX/CypPwzP2ke/Qa2bA=
 github.com/projectdiscovery/utils v0.11.9/go.mod h1:d3BcbsHhuxW9T3jpvrgzwshvqhaL9qPSkoi+3h/DY8E=
+github.com/projectdiscovery/utils v0.11.10-0.20261017183053-ac02d12acdf3 h1:TcGBnYp/RhcGGwNa6sP8/V/oJBDsbGnmBWU1yR0K/2U=
+github.com/projectdiscovery/utils v0.11.10-0.20261017183053-ac02d12acdf3/go.mod h1:d3BcbsHhuxW9T3jpvrgzwshvqhaL9qPSkoi+3h/DY8E=
 github.com/projectdiscovery/wappalyzergo v0.3.4 h1:yO+xel+C0d/vOFWcUKVZRpIDAfcEmsjpVDjQ0jifAUw=
 github.com/projectdiscovery/wappalyzergo v0.3.4/go.mod h1:6xTRCUQhX8m+L1eqJDQIALJlRv2Bnhrmku80YU9bBXE=
 github.com/projectdiscovery/yamldoc-go v1.0.7 h1:Odmjd1o4EkMff2F6QHi1dIYPIQPZUFy24YqmQCprmTk=

pkg/protocols/common/protocolstate/memguardian.go

@@ -5,6 +5,7 @@ import (
 	"sync"
 	"time"
 
+	"github.com/docker/go-units"
 	"github.com/projectdiscovery/gologger"
 	"github.com/projectdiscovery/utils/env"
 	httputil "github.com/projectdiscovery/utils/http"
@@ -16,6 +17,7 @@ var (
 	MaxBytesBufferAllocOnLowMemory = env.GetEnvOrDefault("MEMGUARDIAN_ALLOC", 0)
 	memTimer                       *time.Ticker
 	cancelFunc                     context.CancelFunc
+	memWatcherDone                 chan struct{}
 	muGlobalChange                 sync.Mutex
 )
 
@@ -25,12 +27,18 @@ func StartActiveMemGuardian(ctx context.Context) {
 	if memguardian.DefaultMemGuardian == nil || memTimer != nil {
 		return
 	}
+	if err := refreshMemoryWarning(); err != nil {
+		gologger.Warning().Msgf("could not check memory pressure: %s", err)
+	}
 
 	memTimer = time.NewTicker(memguardian.DefaultInterval)
 	ctx, cancelFunc = context.WithCancel(ctx)
+	memWatcherDone = make(chan struct{})
 
 	ticker := memTimer
+	done := memWatcherDone
 	go func(t *time.Ticker) {
+		defer close(done)
 		if t == nil {
 			return
 		}
@@ -39,6 +47,9 @@ func StartActiveMemGuardian(ctx context.Context) {
 			case <-ctx.Done():
 				return
 			case <-t.C:
+				if err := refreshMemoryWarning(); err != nil {
+					gologger.Warning().Msgf("could not check memory pressure: %s", err)
+				}
 				if IsLowOnMemory() {
 					_ = GlobalGuardBytesBufferAlloc()
 				} else {
@@ -49,6 +60,24 @@ func StartActiveMemGuardian(ctx context.Context) {
 	}(ticker)
 }
 
+func refreshMemoryWarning() error {
+	usedRatio, usedBytes, err := memguardian.UsedRam()
+	if err != nil {
+		return err
+	}
+
+	warning := usedRatio >= memguardian.DefaultMaxUsedRamRatio
+	if maxRam := env.GetEnvOrDefault(memguardian.MemGuardianMaxUsedMemoryENV, ""); maxRam != "" {
+		maxBytes, err := units.FromHumanSize(maxRam)
+		if err != nil {
+			return err
+		}
+		warning = warning || usedBytes >= uint64(maxBytes)
+	}
+	memguardian.DefaultMemGuardian.Warning.Store(warning)
+	return nil
+}
+
 func StopActiveMemGuardian() {
 	muGlobalChange.Lock()
 	defer muGlobalChange.Unlock()
@@ -61,6 +90,10 @@ func StopActiveMemGuardian() {
 		cancelFunc()
 		cancelFunc = nil
 	}
+	if memWatcherDone != nil {
+		<-memWatcherDone
+		memWatcherDone = nil
+	}
 	if memTimer != nil {
 		memTimer.Stop()
 		memTimer = nil

pkg/protocols/common/protocolstate/memguardian_test.go

@@ -127,3 +127,33 @@ func TestMemGuardianReset(t *testing.T) {
 	StopActiveMemGuardian()
 	time.Sleep(10 * time.Millisecond) // Allow cleanup
 }
+
+func TestActiveMemGuardianUpdatesLowMemoryWarning(t *testing.T) {
+	StopActiveMemGuardian()
+	t.Setenv(memguardian.MemGuardianMaxUsedMemoryENV, "")
+	original := memguardian.DefaultMemGuardian
+	originalInterval := memguardian.DefaultInterval
+	originalRatio := memguardian.DefaultMaxUsedRamRatio
+	memguardian.DefaultInterval = 5 * time.Millisecond
+	memguardian.DefaultMaxUsedRamRatio = 1
+	guardian, err := memguardian.New(
+		memguardian.WitInterval(memguardian.DefaultInterval),
+		memguardian.WithMaxRamRatioWarning(1),
+	)
+	require.NoError(t, err)
+	memguardian.DefaultMemGuardian = guardian
+	t.Cleanup(func() {
+		StopActiveMemGuardian()
+		guardian.Close()
+		memguardian.DefaultMemGuardian = original
+		memguardian.DefaultInterval = originalInterval
+		memguardian.DefaultMaxUsedRamRatio = originalRatio
+	})
+
+	StartActiveMemGuardian(context.Background())
+	require.Eventually(t, IsLowOnMemory, 100*time.Millisecond, 5*time.Millisecond)
+	StopActiveMemGuardian()
+	memguardian.DefaultMaxUsedRamRatio = 100
+	require.NoError(t, refreshMemoryWarning())
+	require.False(t, IsLowOnMemory())
+}