Skip to content

Commit

Permalink
Merge pull request #7594 from vmg/vmg/perf-signal
Browse files Browse the repository at this point in the history
pprof: allow stopping profiling early with a signal
  • Loading branch information
deepthi authored Mar 4, 2021
2 parents 2c0c29b + 42456fb commit 434115b
Showing 1 changed file with 33 additions and 14 deletions.
47 changes: 33 additions & 14 deletions go/vt/servenv/pprof.go
Original file line number Diff line number Diff line change
Expand Up @@ -21,13 +21,15 @@ import (
"fmt"
"io/ioutil"
"os"
"os/signal"
"path/filepath"
"runtime"
"runtime/pprof"
"runtime/trace"
"strconv"
"strings"
"sync/atomic"
"syscall"

"vitess.io/vitess/go/vt/log"
)
Expand Down Expand Up @@ -134,6 +136,14 @@ func parseProfileFlag(pf string) (*profile, error) {

var profileStarted uint32

func stopCallback(stop func()) func() {
return func() {
if atomic.CompareAndSwapUint32(&profileStarted, 1, 0) {
stop()
}
}
}

// start begins the configured profiling process and returns a cleanup function
// that must be executed before process termination to flush the profile to disk.
// Based on the profiling code in github.com/pkg/profile
Expand Down Expand Up @@ -172,66 +182,66 @@ func (prof *profile) start() func() {
switch prof.mode {
case profileCPU:
pprof.StartCPUProfile(f)
return func() {
return stopCallback(func() {
pprof.StopCPUProfile()
f.Close()
}
})

case profileMemHeap, profileMemAllocs:
old := runtime.MemProfileRate
runtime.MemProfileRate = prof.rate
return func() {
return stopCallback(func() {
tt := "heap"
if prof.mode == profileMemAllocs {
tt = "allocs"
}
pprof.Lookup(tt).WriteTo(f, 0)
f.Close()
runtime.MemProfileRate = old
}
})

case profileMutex:
runtime.SetMutexProfileFraction(prof.rate)
return func() {
return stopCallback(func() {
if mp := pprof.Lookup("mutex"); mp != nil {
mp.WriteTo(f, 0)
}
f.Close()
runtime.SetMutexProfileFraction(0)
}
})

case profileBlock:
runtime.SetBlockProfileRate(prof.rate)
return func() {
return stopCallback(func() {
pprof.Lookup("block").WriteTo(f, 0)
f.Close()
runtime.SetBlockProfileRate(0)
}
})

case profileThreads:
return func() {
return stopCallback(func() {
if mp := pprof.Lookup("threadcreate"); mp != nil {
mp.WriteTo(f, 0)
}
f.Close()
}
})

case profileTrace:
if err := trace.Start(f); err != nil {
log.Fatalf("pprof: could not start trace: %v", err)
}
return func() {
return stopCallback(func() {
trace.Stop()
f.Close()
}
})

case profileGoroutine:
return func() {
return stopCallback(func() {
if mp := pprof.Lookup("goroutine"); mp != nil {
mp.WriteTo(f, 0)
}
f.Close()
}
})

default:
panic("unsupported profile mode")
Expand All @@ -246,6 +256,15 @@ func init() {
}
if prof != nil {
stop := prof.start()

go func() {
ch := make(chan os.Signal, 1)
signal.Notify(ch, syscall.SIGUSR1)

<-ch
stop()
}()

OnTerm(stop)
}
})
Expand Down

0 comments on commit 434115b

Please sign in to comment.