Merge pull request #6718 from projectdiscovery/dwisiswant0/perf/generators/optimize-MergeMaps-to-reduce-allocs

perf(generators): optimize `MergeMaps` to reduce allocs

Author: Mzack9999

lib/sdk.go

@@ -19,6 +19,7 @@ import (
 	"github.com/projectdiscovery/nuclei/v3/pkg/output"
 	"github.com/projectdiscovery/nuclei/v3/pkg/progress"
 	"github.com/projectdiscovery/nuclei/v3/pkg/protocols"
+	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/generators"
 	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/hosterrorscache"
 	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/interactsh"
 	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/protocolinit"
@@ -238,6 +239,9 @@ func (e *NucleiEngine) closeInternal() {
 	if e.tmpDir != "" {
 		_ = os.RemoveAll(e.tmpDir)
 	}
+	if e.opts != nil {
+		generators.ClearOptionsPayloadMap(e.opts)
+	}
 }
 
 // Close all resources used by nuclei engine

pkg/protocols/common/generators/maps.go

@@ -44,20 +44,51 @@ func MergeMapsMany(maps ...interface{}) map[string][]string {
 	return m
 }
 
-// MergeMaps merges two maps into a new map
+// MergeMaps merges multiple maps into a new map.
+//
+// Use [CopyMap] if you need to copy a single map.
+// Use [MergeMapsInto] to merge into an existing map.
 func MergeMaps(maps ...map[string]interface{}) map[string]interface{} {
-	merged := make(map[string]interface{})
+	mapsLen := 0
+	for _, m := range maps {
+		mapsLen += len(m)
+	}
+
+	merged := make(map[string]interface{}, mapsLen)
 	for _, m := range maps {
 		maps0.Copy(merged, m)
 	}
+
 	return merged
 }
 
+// CopyMap creates a shallow copy of a single map.
+func CopyMap(m map[string]interface{}) map[string]interface{} {
+	if m == nil {
+		return nil
+	}
+
+	result := make(map[string]interface{}, len(m))
+	maps0.Copy(result, m)
+
+	return result
+}
+
+// MergeMapsInto copies all entries from src maps into dst (mutating dst).
+//
+// Use when dst is a fresh map the caller owns and wants to avoid allocation.
+func MergeMapsInto(dst map[string]interface{}, srcs ...map[string]interface{}) {
+	for _, src := range srcs {
+		maps0.Copy(dst, src)
+	}
+}
+
 // ExpandMapValues converts values from flat string to string slice
 func ExpandMapValues(m map[string]string) map[string][]string {
 	m1 := make(map[string][]string, len(m))
 	for k, v := range m {
 		m1[k] = []string{v}
 	}
+
 	return m1
 }

pkg/protocols/common/generators/maps_bench_test.go

@@ -0,0 +1,109 @@
+package generators
+
+import (
+	"fmt"
+	"testing"
+)
+
+func BenchmarkMergeMaps(b *testing.B) {
+	map1 := map[string]interface{}{
+		"key1": "value1",
+		"key2": "value2",
+		"key3": "value3",
+		"key4": "value4",
+		"key5": "value5",
+	}
+	map2 := map[string]interface{}{
+		"key6":  "value6",
+		"key7":  "value7",
+		"key8":  "value8",
+		"key9":  "value9",
+		"key10": "value10",
+	}
+	map3 := map[string]interface{}{
+		"key11": "value11",
+		"key12": "value12",
+		"key13": "value13",
+	}
+
+	for i := 1; i <= 3; i++ {
+		b.Run(fmt.Sprintf("%d-maps", i), func(b *testing.B) {
+			b.ReportAllocs()
+			for b.Loop() {
+				switch i {
+				case 1:
+					_ = MergeMaps(map1)
+				case 2:
+					_ = MergeMaps(map1, map2)
+				case 3:
+					_ = MergeMaps(map1, map2, map3)
+				}
+			}
+		})
+	}
+}
+
+func BenchmarkCopyMap(b *testing.B) {
+	map1 := map[string]interface{}{
+		"key1": "value1",
+		"key2": "value2",
+		"key3": "value3",
+		"key4": "value4",
+		"key5": "value5",
+	}
+
+	for i := 1; i <= 1; i++ {
+		b.Run(fmt.Sprintf("%d-maps", i), func(b *testing.B) {
+			b.ReportAllocs()
+			for b.Loop() {
+				switch i {
+				case 1:
+					_ = CopyMap(map1)
+				}
+			}
+		})
+	}
+}
+
+func BenchmarkMergeMapsInto(b *testing.B) {
+	map1 := map[string]interface{}{
+		"key1": "value1",
+		"key2": "value2",
+		"key3": "value3",
+		"key4": "value4",
+		"key5": "value5",
+	}
+	map2 := map[string]interface{}{
+		"key6":  "value6",
+		"key7":  "value7",
+		"key8":  "value8",
+		"key9":  "value9",
+		"key10": "value10",
+	}
+	map3 := map[string]interface{}{
+		"key11": "value11",
+		"key12": "value12",
+		"key13": "value13",
+	}
+	map4 := map[string]interface{}{
+		"key14": "value14",
+		"key15": "value15",
+		"key16": "value16",
+	}
+
+	for i := 1; i <= 3; i++ {
+		b.Run(fmt.Sprintf("%d-maps", i), func(b *testing.B) {
+			b.ReportAllocs()
+			for b.Loop() {
+				switch i {
+				case 1:
+					MergeMapsInto(map1, map2)
+				case 2:
+					MergeMapsInto(map1, map2, map3)
+				case 3:
+					MergeMapsInto(map1, map2, map3, map4)
+				}
+			}
+		})
+	}
+}

pkg/protocols/common/generators/options.go

@@ -1,12 +1,32 @@
 package generators
 
 import (
+	"sync"
+
 	"github.com/projectdiscovery/nuclei/v3/pkg/types"
 )
 
-// BuildPayloadFromOptions returns a map with the payloads provided via CLI
+// optionsPayloadMap caches the result of BuildPayloadFromOptions per options
+// pointer. This supports multiple SDK instances with different options running
+// concurrently.
+var optionsPayloadMap sync.Map // map[*types.Options]map[string]interface{}
+
+// BuildPayloadFromOptions returns a map with the payloads provided via CLI.
+//
+// The result is cached per options pointer since options don't change during a run.
+// Returns a copy of the cached map to prevent concurrent modification issues.
+// Safe for concurrent use with multiple SDK instances.
 func BuildPayloadFromOptions(options *types.Options) map[string]interface{} {
+	if options == nil {
+		return make(map[string]interface{})
+	}
+
+	if cached, ok := optionsPayloadMap.Load(options); ok {
+		return CopyMap(cached.(map[string]interface{}))
+	}
+
 	m := make(map[string]interface{})
+
 	// merge with vars
 	if !options.Vars.IsEmpty() {
 		m = MergeMaps(m, options.Vars.AsMap())
@@ -16,5 +36,18 @@ func BuildPayloadFromOptions(options *types.Options) map[string]interface{} {
 	if options.EnvironmentVariables {
 		m = MergeMaps(EnvVars(), m)
 	}
-	return m
+
+	actual, _ := optionsPayloadMap.LoadOrStore(options, m)
+
+	// Return a copy to prevent concurrent writes to the cached map
+	return CopyMap(actual.(map[string]interface{}))
+}
+
+// ClearOptionsPayloadMap clears the cached options payload.
+// SDK users should call this when disposing of a NucleiEngine instance
+// to prevent memory leaks if creating many short-lived instances.
+func ClearOptionsPayloadMap(options *types.Options) {
+	if options != nil {
+		optionsPayloadMap.Delete(options)
+	}
 }

pkg/protocols/common/generators/options_bench_test.go

@@ -0,0 +1,45 @@
+package generators
+
+import (
+	"testing"
+
+	"github.com/projectdiscovery/goflags"
+	"github.com/projectdiscovery/nuclei/v3/pkg/types"
+)
+
+func BenchmarkBuildPayloadFromOptions(b *testing.B) {
+	// Setup options with vars and env vars
+	vars := goflags.RuntimeMap{}
+	_ = vars.Set("key1=value1")
+	_ = vars.Set("key2=value2")
+	_ = vars.Set("key3=value3")
+	_ = vars.Set("key4=value4")
+	_ = vars.Set("key5=value5")
+
+	opts := &types.Options{
+		Vars:                 vars,
+		EnvironmentVariables: true, // This adds more entries
+	}
+
+	b.Run("Sequential", func(b *testing.B) {
+		ClearOptionsPayloadMap(opts)
+
+		b.ReportAllocs()
+		for b.Loop() {
+			_ = BuildPayloadFromOptions(opts)
+		}
+	})
+
+	b.Run("Parallel", func(b *testing.B) {
+		ClearOptionsPayloadMap(opts)
+
+		b.ReportAllocs()
+		b.RunParallel(func(pb *testing.PB) {
+			for pb.Next() {
+				m := BuildPayloadFromOptions(opts)
+				// Simulate typical usage - read a value
+				_ = m["key1"]
+			}
+		})
+	})
+}

pkg/protocols/common/generators/options_test.go

@@ -0,0 +1,92 @@
+package generators
+
+import (
+	"sync"
+	"testing"
+
+	"github.com/projectdiscovery/goflags"
+	"github.com/projectdiscovery/nuclei/v3/pkg/types"
+	"github.com/stretchr/testify/require"
+)
+
+func TestBuildPayloadFromOptionsConcurrency(t *testing.T) {
+	// Test that BuildPayloadFromOptions is safe for concurrent use
+	// and returns independent copies that can be modified without races
+	vars := goflags.RuntimeMap{}
+	_ = vars.Set("key=value")
+
+	opts := &types.Options{
+		Vars: vars,
+	}
+
+	const numGoroutines = 100
+	var wg sync.WaitGroup
+	wg.Add(numGoroutines)
+
+	// Each goroutine gets a map and modifies it
+	for i := 0; i < numGoroutines; i++ {
+		go func(id int) {
+			defer wg.Done()
+
+			// Get the map (should be a copy of cached data)
+			m := BuildPayloadFromOptions(opts)
+
+			// Modify it - this should not cause races
+			m["goroutine_id"] = id
+			m["test_key"] = "test_value"
+
+			// Verify original cached value is present
+			require.Equal(t, "value", m["key"])
+		}(i)
+	}
+
+	wg.Wait()
+}
+
+func TestBuildPayloadFromOptionsCaching(t *testing.T) {
+	// Test that caching actually works
+	vars := goflags.RuntimeMap{}
+	_ = vars.Set("cached=yes")
+
+	opts := &types.Options{
+		Vars:                 vars,
+		EnvironmentVariables: false,
+	}
+
+	// First call - builds and caches
+	m1 := BuildPayloadFromOptions(opts)
+	require.Equal(t, "yes", m1["cached"])
+
+	// Second call - should return copy of cached result
+	m2 := BuildPayloadFromOptions(opts)
+	require.Equal(t, "yes", m2["cached"])
+
+	// Modify m1 - should not affect m2 since they're copies
+	m1["modified"] = "in_m1"
+	require.NotContains(t, m2, "modified")
+
+	// Modify m2 - should not affect future calls
+	m2["modified"] = "in_m2"
+	m3 := BuildPayloadFromOptions(opts)
+	require.NotContains(t, m3, "modified")
+}
+
+func TestClearOptionsPayloadMap(t *testing.T) {
+	vars := goflags.RuntimeMap{}
+	_ = vars.Set("temp=data")
+
+	opts := &types.Options{
+		Vars: vars,
+	}
+
+	// Build and cache
+	m1 := BuildPayloadFromOptions(opts)
+	require.Equal(t, "data", m1["temp"])
+
+	// Clear the cache
+	ClearOptionsPayloadMap(opts)
+
+	// Verify it still works (rebuilds)
+	m2 := BuildPayloadFromOptions(opts)
+	require.Equal(t, "data", m2["temp"])
+}