Skip to content

Commit 3cb03c0

Browse files
authored
Merge pull request #574 from projectdiscovery/pprof-server-add
feat: added pprof server mode addition
2 parents 70cf636 + b147e26 commit 3cb03c0

File tree

5 files changed

+177
-59
lines changed

5 files changed

+177
-59
lines changed

go.mod

+2
Original file line numberDiff line numberDiff line change
@@ -56,9 +56,11 @@ require (
5656
github.com/dimchansky/utfbom v1.1.1 // indirect
5757
github.com/dlclark/regexp2 v1.11.4 // indirect
5858
github.com/fatih/color v1.15.0 // indirect
59+
github.com/felixge/fgprof v0.9.5 // indirect
5960
github.com/gaissmai/bart v0.9.5 // indirect
6061
github.com/go-ole/go-ole v1.2.6 // indirect
6162
github.com/golang/snappy v0.0.4 // indirect
63+
github.com/google/pprof v0.0.0-20240227163752-401108e1b7e7 // indirect
6264
github.com/json-iterator/go v1.1.12 // indirect
6365
github.com/klauspost/pgzip v1.2.5 // indirect
6466
github.com/lucasb-eyer/go-colorful v1.2.0 // indirect

go.sum

+19
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,12 @@ github.com/charmbracelet/x/exp/golden v0.0.0-20240806155701-69247e0abc2a h1:G99k
4747
github.com/charmbracelet/x/exp/golden v0.0.0-20240806155701-69247e0abc2a/go.mod h1:wDlXFlCrmJ8J+swcL/MnGUuYnqgQdW9rhSD61oNMb6U=
4848
github.com/cheggaaa/pb/v3 v3.1.4 h1:DN8j4TVVdKu3WxVwcRKu0sG00IIU6FewoABZzXbRQeo=
4949
github.com/cheggaaa/pb/v3 v3.1.4/go.mod h1:6wVjILNBaXMs8c21qRiaUM8BR82erfgau1DQ4iUXmSA=
50+
github.com/chromedp/cdproto v0.0.0-20230802225258-3cf4e6d46a89/go.mod h1:GKljq0VrfU4D5yc+2qA6OVr8pmO/MBbPEWqWQ/oqGEs=
51+
github.com/chromedp/chromedp v0.9.2/go.mod h1:LkSXJKONWTCHAfQasKFUZI+mxqS4tZqhmtGzzhLsnLs=
52+
github.com/chromedp/sysutil v1.0.0/go.mod h1:kgWmDdq8fTzXYcKIBqIYvRRTnYb9aNS9moAV0xufSww=
53+
github.com/chzyer/logex v1.2.1/go.mod h1:JLbx6lG2kDbNRFnfkgvh4eRJRPX1QCoOIWomwysCBrQ=
54+
github.com/chzyer/readline v1.5.1/go.mod h1:Eh+b79XXUwfKfcPLepksvw2tcLE/Ct21YObkaSkeBlk=
55+
github.com/chzyer/test v1.0.0/go.mod h1:2JlltgoNkt4TW/z9V/IzDdFaMTM2JPIi26O1pF38GC8=
5056
github.com/cloudflare/circl v1.3.7 h1:qlCDlTPz2n9fu58M0Nh1J/JzcFpfgkFHHX3O35r5vcU=
5157
github.com/cloudflare/circl v1.3.7/go.mod h1:sRTcRWXGLrKw6yIGJ+l7amYJFfAXbZG0kBSc8r4zxgA=
5258
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
@@ -65,6 +71,8 @@ github.com/ebitengine/purego v0.4.0 h1:RQVuMIxQPQ5iCGEJvjQ17YOK+1tMKjVau2FUMvXH4
6571
github.com/ebitengine/purego v0.4.0/go.mod h1:ah1In8AOtksoNK6yk5z1HTJeUkC1Ez4Wk2idgGslMwQ=
6672
github.com/fatih/color v1.15.0 h1:kOqh6YHBtK8aywxGerMG2Eq3H6Qgoqeo13Bk2Mv/nBs=
6773
github.com/fatih/color v1.15.0/go.mod h1:0h5ZqXfHYED7Bhv2ZJamyIOUej9KtShiJESRwBDUSsw=
74+
github.com/felixge/fgprof v0.9.5 h1:8+vR6yu2vvSKn08urWyEuxx75NWPEvybbkBirEpsbVY=
75+
github.com/felixge/fgprof v0.9.5/go.mod h1:yKl+ERSa++RYOs32d8K6WEXCB4uXdLls4ZaZPpayhMM=
6876
github.com/fortytw2/leaktest v1.3.0 h1:u8491cBMTQ8ft8aeV+adlcytMZylmA5nnwwkRZjI8vw=
6977
github.com/fortytw2/leaktest v1.3.0/go.mod h1:jDsjWgpAGjm2CA7WthBh/CdZYEPF31XHquHwclZch5g=
7078
github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo=
@@ -74,6 +82,9 @@ github.com/gaissmai/bart v0.9.5 h1:vy+r4Px6bjZ+v2QYXAsg63vpz9IfzdW146A8Cn4GPIo=
7482
github.com/gaissmai/bart v0.9.5/go.mod h1:KHeYECXQiBjTzQz/om2tqn3sZF1J7hw9m6z41ftj3fg=
7583
github.com/go-ole/go-ole v1.2.6 h1:/Fpf6oFPoeFik9ty7siob0G6Ke8QvQEuVcuChpwXzpY=
7684
github.com/go-ole/go-ole v1.2.6/go.mod h1:pprOEPIfldk/42T2oK7lQ4v4JSDwmV0As9GaiUsvbm0=
85+
github.com/gobwas/httphead v0.1.0/go.mod h1:O/RXo79gxV8G+RqlR/otEwx4Q36zl9rqC5u12GKvMCM=
86+
github.com/gobwas/pool v0.2.1/go.mod h1:q8bcK0KcYlCgd9e7WYLm9LpyS+YeLd8JVDW6WezmKEw=
87+
github.com/gobwas/ws v1.2.1/go.mod h1:hRKAFb8wOxFROYNsT1bqfWnhX+b5MFeJM9r2ZSwg/KY=
7788
github.com/gofrs/uuid v3.3.0+incompatible/go.mod h1:b2aQJv3Z4Fp6yNu3cdSllBxTCLRxnplIgP/c0N/04lM=
7889
github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
7990
github.com/golang/protobuf v1.3.1/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
@@ -100,6 +111,8 @@ github.com/google/go-querystring v1.0.0/go.mod h1:odCYkC5MyYFN7vkCjXpyrEuKhc/BUO
100111
github.com/google/go-querystring v1.1.0 h1:AnCroh3fv4ZBgVIf1Iwtovgjaw/GiKJo8M8yD/fhyJ8=
101112
github.com/google/go-querystring v1.1.0/go.mod h1:Kcdr2DB4koayq7X8pmAG4sNG59So17icRSOU623lUBU=
102113
github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
114+
github.com/google/pprof v0.0.0-20240227163752-401108e1b7e7 h1:y3N7Bm7Y9/CtpiVkw/ZWj6lSlDF3F74SfKwfTCer72Q=
115+
github.com/google/pprof v0.0.0-20240227163752-401108e1b7e7/go.mod h1:czg5+yv1E0ZGTi6S6vVK1mke0fV+FaUhNGcd6VRS9Ik=
103116
github.com/google/uuid v1.3.1 h1:KjJaJ9iWZ3jOFZIf1Lqf4laDRCasjl0BCmnEGxkdLb4=
104117
github.com/google/uuid v1.3.1/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
105118
github.com/gorilla/css v1.0.1 h1:ntNaBIghp6JmvWnxbZKANoLyuXTPZ4cAMlo6RyhlbO8=
@@ -111,6 +124,8 @@ github.com/hdm/jarm-go v0.0.7/go.mod h1:kinGoS0+Sdn1Rr54OtanET5E5n7AlD6T6CrJAKDj
111124
github.com/hexops/gotextdiff v1.0.3 h1:gitA9+qJrrTCsiCl7+kh75nPqQt1cx4ZkudSTLoUqJM=
112125
github.com/hexops/gotextdiff v1.0.3/go.mod h1:pSWU5MAI3yDq+fZBTazCSJysOMbxWL1BSow5/V2vxeg=
113126
github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU=
127+
github.com/ianlancetaylor/demangle v0.0.0-20230524184225-eabc099b10ab/go.mod h1:gx7rwoVhcfuVKG5uya9Hs3Sxj7EIvldVofAWIUtGouw=
128+
github.com/josharian/intern v1.0.0/go.mod h1:5DoeVV0s6jJacbCEi61lwdGj/aVlrQvzHFFd8Hwg//Y=
114129
github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM=
115130
github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo=
116131
github.com/julienschmidt/httprouter v1.3.0 h1:U0609e9tgbseu3rBINet9P48AI/D3oJs4dN7jwJOQ1U=
@@ -131,12 +146,14 @@ github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
131146
github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
132147
github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
133148
github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
149+
github.com/ledongthuc/pdf v0.0.0-20220302134840-0c2507a12d80/go.mod h1:imJHygn/1yfhB7XSJJKlFZKl/J+dCPAknuiaGOshXAs=
134150
github.com/logrusorgru/aurora v2.0.3+incompatible h1:tOpm7WcpBTn4fjmVfgpQq0EfczGlG91VSDkswnjF5A8=
135151
github.com/logrusorgru/aurora v2.0.3+incompatible/go.mod h1:7rIyQOR62GCctdiQpZ/zOJlFyk6y+94wXzv6RNZgaR4=
136152
github.com/lucasb-eyer/go-colorful v1.2.0 h1:1nnpGOrhyZZuNyfu1QjKiUICQ74+3FNCN69Aj6K7nkY=
137153
github.com/lucasb-eyer/go-colorful v1.2.0/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0=
138154
github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0 h1:6E+4a0GO5zZEnZ81pIr0yLvtUWk2if982qA3F3QD6H4=
139155
github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0/go.mod h1:zJYVVT2jmtg6P3p1VtQj7WsuWi/y4VnjVBn7F8KPB3I=
156+
github.com/mailru/easyjson v0.7.7/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJJLY9Nlc=
140157
github.com/mattn/go-colorable v0.1.13 h1:fFA4WZxdEF4tXPZVKMLwD8oUnCTTo08duU7wxecdEvA=
141158
github.com/mattn/go-colorable v0.1.13/go.mod h1:7S9/ev0klgBDR4GtXTXX8a3vIGJpMovkB8vQcUbaXHg=
142159
github.com/mattn/go-isatty v0.0.16/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM=
@@ -178,6 +195,7 @@ github.com/onsi/gomega v1.4.3/go.mod h1:ex+gbHU/CVuBBDIJjb2X0qEXbFg53c61hWP/1Cpa
178195
github.com/onsi/gomega v1.27.6 h1:ENqfyGeS5AX/rlXDd/ETokDz93u0YufY1Pgxuy/PvWE=
179196
github.com/onsi/gomega v1.27.6/go.mod h1:PIQNjfQwkP3aQAH7lf7j87O/5FiNr+ZR8+ipb+qQlhg=
180197
github.com/op/go-logging v0.0.0-20160315200505-970db520ece7/go.mod h1:HzydrMdWErDVzsI23lYNej1Htcns9BCg93Dk0bBINWk=
198+
github.com/orisano/pixelmatch v0.0.0-20220722002657-fb0b55479cde/go.mod h1:nZgzbfBr3hhjoZnS66nKrHmduYNpc34ny7RK4z5/HM0=
181199
github.com/pierrec/lz4/v4 v4.1.2 h1:qvY3YFXRQE/XB8MlLzJH7mSzBs74eA2gg52YTk6jUPM=
182200
github.com/pierrec/lz4/v4 v4.1.2/go.mod h1:gZWDp/Ze/IJXGXf23ltt2EXimqmTUXEy0GFuRQyBid4=
183201
github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
@@ -370,6 +388,7 @@ golang.org/x/sys v0.0.0-20201204225414-ed752295db88/go.mod h1:h1NjWce9XRLGQEsW7w
370388
golang.org/x/sys v0.0.0-20210228012217-479acdf4ea46/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
371389
golang.org/x/sys v0.0.0-20210423082822-04245dca01da/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
372390
golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
391+
golang.org/x/sys v0.0.0-20220310020820-b874c991c1a5/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
373392
golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
374393
golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
375394
golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=

pprof/README.md

+36-5
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,13 @@
11
## PProfiling Usage Guide
22

3-
### Environment Variables
3+
Two types of profiling are supported:
4+
5+
1. **pprof**: Standard go profiling writing to files in a directory.
6+
2. **server**: Profiling server listening on a port with pprof and fgprof endpoints.
7+
8+
### pprof
9+
10+
#### Environment Variables
411

512
- `PPROF`: Enable or disable profiling. Set to 1 to enable.
613
- `MEM_PROFILE_DIR`: Directory to write memory profiles to.
@@ -9,7 +16,7 @@
916
- `MEM_PROFILE_RATE`: Memory profiling rate (default 4096).
1017

1118

12-
## How to Use
19+
### How to Use
1320

1421
1. Set the environment variables as per your requirements.
1522

@@ -30,7 +37,7 @@ export MEM_PROFILE_RATE=4096
3037
- Profiles will be written at intervals specified by PPROF_TIME.
3138
- Memory profiling rate is controlled by MEM_PROFILE_RATE.
3239

33-
### Example
40+
#### Example
3441

3542
```bash
3643
[+] GOOS: linux
@@ -46,7 +53,31 @@ profile: ticker enabled (rate 10s)
4653
profile: cpu profiling enabled (ticker 10s)
4754
```
4855

49-
### Note
56+
#### Note
5057

5158
- The polling time (PPROF_TIME) should be set according to your application's performance and profiling needs.
52-
- The memory profiling rate (MEM_PROFILE_RATE) controls the granularity of the memory profiling. Higher values provide more detail but consume more resources.
59+
- The memory profiling rate (MEM_PROFILE_RATE) controls the granularity of the memory profiling. Higher values provide more detail but consume more resources.
60+
61+
### server
62+
63+
Server is a simple HTTP server listening on a port with pprof and fgprof endpoints.
64+
65+
#### Environment Variables
66+
67+
- `PPROF_SERVER_ADDR`: Address to listen on for pprof and fgprof server (default 127.0.0.1:6060).
68+
69+
#### Endpoints
70+
71+
- /debug/pprof/
72+
- /debug/pprof/cmdline
73+
- /debug/pprof/profile
74+
- /debug/pprof/profile
75+
- /debug/pprof/symbol
76+
- /debug/pprof/trace
77+
- /debug/fgprof
78+
79+
#### Example
80+
81+
```console
82+
go tool pprof http://127.0.0.1:8086/debug/fgprof
83+
```

pprof/pprof.go

+58-54
Original file line numberDiff line numberDiff line change
@@ -24,65 +24,69 @@ const (
2424

2525
func init() {
2626
if env.GetEnvOrDefault(PPROFSwitchENV, 0) == 1 {
27-
log.Printf("[+] GOOS: %v\n", runtime.GOOS)
28-
log.Printf("[+] GOARCH: %v\n", runtime.GOARCH)
29-
log.Printf("[+] Command: %v\n", strings.Join(os.Args, " "))
30-
log.Println("Available PPROF Config Options:")
31-
log.Printf("%-16v - directory to write memory profiles to\n", MemProfileENV)
32-
log.Printf("%-16v - directory to write cpu profiles to\n", CPUProfileENV)
33-
log.Printf("%-16v - polling time for cpu and memory profiles (with unit ex: 10s)\n", PPROFTimeENV)
34-
log.Printf("%-16v - memory profiling rate (default 4096)\n", MemProfileRate)
27+
startDefaultProfiler()
28+
}
29+
}
30+
31+
func startDefaultProfiler() {
32+
log.Printf("[+] GOOS: %v\n", runtime.GOOS)
33+
log.Printf("[+] GOARCH: %v\n", runtime.GOARCH)
34+
log.Printf("[+] Command: %v\n", strings.Join(os.Args, " "))
35+
log.Println("Available PPROF Config Options:")
36+
log.Printf("%-16v - directory to write memory profiles to\n", MemProfileENV)
37+
log.Printf("%-16v - directory to write cpu profiles to\n", CPUProfileENV)
38+
log.Printf("%-16v - polling time for cpu and memory profiles (with unit ex: 10s)\n", PPROFTimeENV)
39+
log.Printf("%-16v - memory profiling rate (default 4096)\n", MemProfileRate)
3540

36-
memProfilesDir := env.GetEnvOrDefault(MemProfileENV, "memdump")
37-
cpuProfilesDir := env.GetEnvOrDefault(CPUProfileENV, "cpuprofile")
38-
pprofTimeDuration := env.GetEnvOrDefault(PPROFTimeENV, time.Duration(3)*time.Second)
39-
pprofRate := env.GetEnvOrDefault(MemProfileRate, 4096)
41+
memProfilesDir := env.GetEnvOrDefault(MemProfileENV, "memdump")
42+
cpuProfilesDir := env.GetEnvOrDefault(CPUProfileENV, "cpuprofile")
43+
pprofTimeDuration := env.GetEnvOrDefault(PPROFTimeENV, time.Duration(3)*time.Second)
44+
pprofRate := env.GetEnvOrDefault(MemProfileRate, 4096)
4045

41-
_ = os.MkdirAll(memProfilesDir, 0755)
42-
_ = os.MkdirAll(cpuProfilesDir, 0755)
46+
_ = os.MkdirAll(memProfilesDir, 0755)
47+
_ = os.MkdirAll(cpuProfilesDir, 0755)
4348

44-
runtime.MemProfileRate = pprofRate
45-
log.Printf("profile: memory profiling enabled (rate %d), %s\n", runtime.MemProfileRate, memProfilesDir)
46-
log.Printf("profile: ticker enabled (rate %s)\n", pprofTimeDuration)
49+
runtime.MemProfileRate = pprofRate
50+
log.Printf("profile: memory profiling enabled (rate %d), %s\n", runtime.MemProfileRate, memProfilesDir)
51+
log.Printf("profile: ticker enabled (rate %s)\n", pprofTimeDuration)
4752

48-
// cpu ticker and profiler
49-
go func() {
50-
ticker := time.NewTicker(pprofTimeDuration)
51-
count := 0
52-
buff := bytes.Buffer{}
53-
log.Printf("profile: cpu profiling enabled (ticker %s)\n", pprofTimeDuration)
54-
for {
55-
err := pprof.StartCPUProfile(&buff)
56-
if err != nil {
57-
log.Fatalf("profile: could not start cpu profile: %s\n", err)
58-
}
59-
<-ticker.C
60-
pprof.StopCPUProfile()
61-
if err := os.WriteFile(filepath.Join(cpuProfilesDir, "cpuprofile-t"+strconv.Itoa(count)+".out"), buff.Bytes(), 0755); err != nil {
62-
log.Fatalf("profile: could not write cpu profile: %s\n", err)
63-
}
64-
buff.Reset()
65-
count++
53+
// cpu ticker and profiler
54+
go func() {
55+
ticker := time.NewTicker(pprofTimeDuration)
56+
count := 0
57+
buff := bytes.Buffer{}
58+
log.Printf("profile: cpu profiling enabled (ticker %s)\n", pprofTimeDuration)
59+
for {
60+
err := pprof.StartCPUProfile(&buff)
61+
if err != nil {
62+
log.Fatalf("profile: could not start cpu profile: %s\n", err)
6663
}
67-
}()
64+
<-ticker.C
65+
pprof.StopCPUProfile()
66+
if err := os.WriteFile(filepath.Join(cpuProfilesDir, "cpuprofile-t"+strconv.Itoa(count)+".out"), buff.Bytes(), 0755); err != nil {
67+
log.Fatalf("profile: could not write cpu profile: %s\n", err)
68+
}
69+
buff.Reset()
70+
count++
71+
}
72+
}()
6873

69-
// memory ticker and profiler
70-
go func() {
71-
ticker := time.NewTicker(pprofTimeDuration)
72-
count := 0
73-
log.Printf("profile: memory profiling enabled (ticker %s)\n", pprofTimeDuration)
74-
for {
75-
<-ticker.C
76-
var buff bytes.Buffer
77-
if err := pprof.WriteHeapProfile(&buff); err != nil {
78-
log.Printf("profile: could not write memory profile: %s\n", err)
79-
}
80-
err := os.WriteFile(filepath.ToSlash(filepath.Join(memProfilesDir, "memprofile-t"+strconv.Itoa(count)+".out")), buff.Bytes(), 0755)
81-
if err != nil {
82-
log.Printf("profile: could not write memory profile: %s\n", err)
83-
}
84-
count++
74+
// memory ticker and profiler
75+
go func() {
76+
ticker := time.NewTicker(pprofTimeDuration)
77+
count := 0
78+
log.Printf("profile: memory profiling enabled (ticker %s)\n", pprofTimeDuration)
79+
for {
80+
<-ticker.C
81+
var buff bytes.Buffer
82+
if err := pprof.WriteHeapProfile(&buff); err != nil {
83+
log.Printf("profile: could not write memory profile: %s\n", err)
8584
}
86-
}()
87-
}
85+
err := os.WriteFile(filepath.ToSlash(filepath.Join(memProfilesDir, "memprofile-t"+strconv.Itoa(count)+".out")), buff.Bytes(), 0755)
86+
if err != nil {
87+
log.Printf("profile: could not write memory profile: %s\n", err)
88+
}
89+
count++
90+
}
91+
}()
8892
}

pprof/server.go

+62
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
package pprof
2+
3+
import (
4+
"context"
5+
"net/http"
6+
"net/http/pprof"
7+
"runtime"
8+
"time"
9+
10+
"github.com/felixge/fgprof"
11+
"github.com/projectdiscovery/gologger"
12+
"github.com/projectdiscovery/utils/env"
13+
)
14+
15+
const (
16+
PPROFServerAddressENV = "PPROF_SERVER_ADDRESS"
17+
)
18+
19+
type PprofServer struct {
20+
server *http.Server
21+
}
22+
23+
func NewPprofServer() *PprofServer {
24+
address := env.GetEnvOrDefault(PPROFServerAddressENV, "127.0.0.1:8086")
25+
26+
mux := http.NewServeMux()
27+
// Default pprof handlers
28+
mux.HandleFunc("/debug/pprof/", pprof.Index)
29+
mux.HandleFunc("/debug/pprof/cmdline", pprof.Cmdline)
30+
mux.HandleFunc("/debug/pprof/profile", pprof.Profile)
31+
mux.HandleFunc("/debug/pprof/symbol", pprof.Symbol)
32+
mux.HandleFunc("/debug/pprof/trace", pprof.Trace)
33+
34+
// Also add fgprof for more detailed profiling
35+
mux.Handle("/debug/fgprof", fgprof.Handler())
36+
37+
server := &http.Server{
38+
Addr: address,
39+
Handler: mux,
40+
}
41+
// Enable block and mutex profiling as well
42+
runtime.SetBlockProfileRate(1)
43+
runtime.SetMutexProfileFraction(1)
44+
45+
return &PprofServer{server: server}
46+
}
47+
48+
func (p *PprofServer) Start() {
49+
gologger.Info().Msgf("Listening pprof debug server on: %s", p.server.Addr)
50+
51+
go func() {
52+
if err := p.server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
53+
gologger.Error().Msgf("pprof server failed to start: %s", err)
54+
}
55+
}()
56+
}
57+
58+
func (p *PprofServer) Stop() {
59+
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
60+
defer cancel()
61+
_ = p.server.Shutdown(ctx)
62+
}

0 commit comments

Comments
 (0)