Merge pull request #266 from projectdiscovery/issue-265-pprof-utils

add pprof utils for performance testing/debugging

env/env.go

@@ -2,7 +2,9 @@ package env
 
 import (
 	"os"
+	"strconv"
 	"strings"
+	"time"
 )
 
 var (
@@ -20,3 +22,46 @@ func ExpandWithEnv(variables ...*string) {
 		*variable = os.Getenv(strings.TrimPrefix(*variable, "$"))
 	}
 }
+
+// EnvType is a type that can be used as a type for environment variables.
+type EnvType interface {
+	~string | ~int | ~bool | ~float64 | time.Duration
+}
+
+// GetEnvOrDefault returns the value of the environment variable or the default value if the variable is not set.
+// in requested type.
+func GetEnvOrDefault[T EnvType](key string, defaultValue T) T {
+	value := os.Getenv(key)
+	if value == "" {
+		return defaultValue
+	}
+	switch any(defaultValue).(type) {
+	case string:
+		return any(value).(T)
+	case int:
+		intVal, err := strconv.Atoi(value)
+		if err != nil || value == "" {
+			return defaultValue
+		}
+		return any(intVal).(T)
+	case bool:
+		boolVal, err := strconv.ParseBool(value)
+		if err != nil || value == "" {
+			return defaultValue
+		}
+		return any(boolVal).(T)
+	case float64:
+		floatVal, err := strconv.ParseFloat(value, 64)
+		if err != nil || value == "" {
+			return defaultValue
+		}
+		return any(floatVal).(T)
+	case time.Duration:
+		durationVal, err := time.ParseDuration(value)
+		if err != nil || value == "" {
+			return defaultValue
+		}
+		return any(durationVal).(T)
+	}
+	return defaultValue
+}

env/env_test.go

@@ -3,6 +3,7 @@ package env
 import (
 	"os"
 	"testing"
+	"time"
 )
 
 func TestExpandWithEnv(t *testing.T) {
@@ -42,3 +43,46 @@ func TestExpandWithEnvNilInput(t *testing.T) {
 	var nilVar *string = nil
 	ExpandWithEnv(nilVar)
 }
+
+func TestGetEnvOrDefault(t *testing.T) {
+	// Test for string
+	os.Setenv("TEST_STRING", "test")
+	resultString := GetEnvOrDefault("TEST_STRING", "default")
+	if resultString != "test" {
+		t.Errorf("Expected 'test', got %s", resultString)
+	}
+
+	// Test for int
+	os.Setenv("TEST_INT", "123")
+	resultInt := GetEnvOrDefault("TEST_INT", 0)
+	if resultInt != 123 {
+		t.Errorf("Expected 123, got %d", resultInt)
+	}
+
+	// Test for bool
+	os.Setenv("TEST_BOOL", "true")
+	resultBool := GetEnvOrDefault("TEST_BOOL", false)
+	if resultBool != true {
+		t.Errorf("Expected true, got %t", resultBool)
+	}
+
+	// Test for float64
+	os.Setenv("TEST_FLOAT", "1.23")
+	resultFloat := GetEnvOrDefault("TEST_FLOAT", 0.0)
+	if resultFloat != 1.23 {
+		t.Errorf("Expected 1.23, got %f", resultFloat)
+	}
+
+	// Test for time.Duration
+	os.Setenv("TEST_DURATION", "1h")
+	resultDuration := GetEnvOrDefault("TEST_DURATION", time.Duration(0))
+	if resultDuration != time.Hour {
+		t.Errorf("Expected 1h, got %s", resultDuration)
+	}
+
+	// Test for default value
+	resultDefault := GetEnvOrDefault("NON_EXISTING", "default")
+	if resultDefault != "default" {
+		t.Errorf("Expected 'default', got %s", resultDefault)
+	}
+}

pprof/README.md

@@ -0,0 +1,52 @@
+## PProfiling Usage Guide
+
+### Environment Variables
+
+- `PPROF`: Enable or disable profiling. Set to 1 to enable.
+- `MEM_PROFILE_DIR`: Directory to write memory profiles to.
+- `CPU_PROFILE_DIR`: Directory to write CPU profiles to.
+- `PPROF_TIME`: Polling time for CPU and memory profiles (with unit ex: 10s).
+- `MEM_PROFILE_RATE`: Memory profiling rate (default 4096).
+
+
+## How to Use
+
+1. Set the environment variables as per your requirements.
+
+```bash
+export PPROF=1
+export MEM_PROFILE_DIR=/path/to/memprofile
+export CPU_PROFILE_DIR=/path/to/cpuprofile
+export PPROF_TIME=10s
+export MEM_PROFILE_RATE=4096
+```
+
+2. Run your Go application. The profiler will start automatically if PPROF is set to 1.
+
+**Output**
+
+- Memory profiles will be written to the directory specified by MEM_PROFILE_DIR.
+- CPU profiles will be written to the directory specified by CPU_PROFILE_DIR.
+- Profiles will be written at intervals specified by PPROF_TIME.
+- Memory profiling rate is controlled by MEM_PROFILE_RATE.
+
+### Example
+
+```bash
+[+] GOOS: linux
+[+] GOARCH: amd64
+[+] Command: /path/to/your/app
+Available PPROF Config Options:
+MEM_PROFILE_DIR   - directory to write memory profiles to
+CPU_PROFILE_DIR   - directory to write cpu profiles to
+PPROF_TIME        - polling time for cpu and memory profiles (with unit ex: 10s)
+MEM_PROFILE_RATE  - memory profiling rate (default 4096)
+profile: memory profiling enabled (rate 4096), /path/to/memprofile
+profile: ticker enabled (rate 10s)
+profile: cpu profiling enabled (ticker 10s)
+```
+
+### Note
+
+- The polling time (PPROF_TIME) should be set according to your application's performance and profiling needs.
+- The memory profiling rate (MEM_PROFILE_RATE) controls the granularity of the memory profiling. Higher values provide more detail but consume more resources.

pprof/pprof.go

@@ -0,0 +1,88 @@
+package pprof
+
+import (
+	"bytes"
+	"log"
+	"os"
+	"path/filepath"
+	"runtime"
+	"runtime/pprof"
+	"strconv"
+	"strings"
+	"time"
+
+	"github.com/projectdiscovery/utils/env"
+)
+
+const (
+	PPROFSwitchENV = "PPROF"
+	MemProfileENV  = "MEM_PROFILE_DIR"
+	CPUProfileENV  = "CPU_PROFILE_DIR"
+	PPROFTimeENV   = "PPROF_TIME"
+	MemProfileRate = "MEM_PROFILE_RATE"
+)
+
+func init() {
+	if env.GetEnvOrDefault(PPROFSwitchENV, 0) == 1 {
+		log.Printf("[+] GOOS: %v\n", runtime.GOOS)
+		log.Printf("[+] GOARCH: %v\n", runtime.GOARCH)
+		log.Printf("[+] Command: %v\n", strings.Join(os.Args, " "))
+		log.Println("Available PPROF Config Options:")
+		log.Printf("%-16v - directory to write memory profiles to\n", MemProfileENV)
+		log.Printf("%-16v - directory to write cpu profiles to\n", CPUProfileENV)
+		log.Printf("%-16v - polling time for cpu and memory profiles (with unit ex: 10s)\n", PPROFTimeENV)
+		log.Printf("%-16v - memory profiling rate (default 4096)\n", MemProfileRate)
+
+		memProfilesDir := env.GetEnvOrDefault(MemProfileENV, "memdump")
+		cpuProfilesDir := env.GetEnvOrDefault(CPUProfileENV, "cpuprofile")
+		pprofTimeDuration := env.GetEnvOrDefault(PPROFTimeENV, time.Duration(3)*time.Second)
+		pprofRate := env.GetEnvOrDefault(MemProfileRate, 4096)
+
+		_ = os.MkdirAll(memProfilesDir, 0755)
+		_ = os.MkdirAll(cpuProfilesDir, 0755)
+
+		runtime.MemProfileRate = pprofRate
+		log.Printf("profile: memory profiling enabled (rate %d), %s\n", runtime.MemProfileRate, memProfilesDir)
+		log.Printf("profile: ticker enabled (rate %s)\n", pprofTimeDuration)
+
+		// cpu ticker and profiler
+		go func() {
+			ticker := time.NewTicker(pprofTimeDuration)
+			count := 0
+			buff := bytes.Buffer{}
+			log.Printf("profile: cpu profiling enabled (ticker %s)\n", pprofTimeDuration)
+			for {
+				err := pprof.StartCPUProfile(&buff)
+				if err != nil {
+					log.Fatalf("profile: could not start cpu profile: %s\n", err)
+				}
+				<-ticker.C
+				pprof.StopCPUProfile()
+				if err := os.WriteFile(filepath.Join(cpuProfilesDir, "cpuprofile-t"+strconv.Itoa(count)+".out"), buff.Bytes(), 0755); err != nil {
+					log.Fatalf("profile: could not write cpu profile: %s\n", err)
+				}
+				buff.Reset()
+				count++
+			}
+		}()
+
+		// memory ticker and profiler
+		go func() {
+			ticker := time.NewTicker(pprofTimeDuration)
+			count := 0
+			log.Printf("profile: memory profiling enabled (ticker %s)\n", pprofTimeDuration)
+			for {
+				<-ticker.C
+				var buff bytes.Buffer
+				if err := pprof.WriteHeapProfile(&buff); err != nil {
+					log.Printf("profile: could not write memory profile: %s\n", err)
+				}
+				err := os.WriteFile(filepath.ToSlash(filepath.Join(memProfilesDir, "memprofile-t"+strconv.Itoa(count)+".out")), buff.Bytes(), 0755)
+				if err != nil {
+					log.Printf("profile: could not write memory profile: %s\n", err)
+				}
+				count++
+			}
+		}()
+	}
+}