new: cpuprofile and memprofile options for easy profiling

This commit is contained in:
evilsocket 2018-02-11 18:41:49 +01:00
commit 47230567ba
4 changed files with 34 additions and 1 deletions

View file

@ -6,6 +6,8 @@ import (
"fmt"
"os"
"os/signal"
"runtime"
"runtime/pprof"
"sort"
"strings"
"syscall"
@ -79,6 +81,14 @@ func New() (*Session, error) {
return nil, err
}
if *s.Options.CpuProfile != "" {
if f, err := os.Create(*s.Options.CpuProfile); err != nil {
return nil, err
} else if err := pprof.StartCPUProfile(f); err != nil {
return nil, err
}
}
s.Env = NewEnvironment(s)
s.Events = NewEventPool(*s.Options.Debug, *s.Options.Silent)
@ -180,6 +190,23 @@ func (s *Session) Close() {
s.Firewall.Restore()
s.Queue.Stop()
if *s.Options.CpuProfile != "" {
pprof.StopCPUProfile()
}
if *s.Options.MemProfile != "" {
f, err := os.Create(*s.Options.MemProfile)
if err != nil {
fmt.Printf("Could not create memory profile: %s\n", err)
return
}
defer f.Close()
runtime.GC() // get up-to-date statistics
if err := pprof.WriteHeapProfile(f); err != nil {
fmt.Printf("Could not write memory profile: %s\n", err)
}
}
}
func (s *Session) Register(mod Module) error {