diff --git a/_examples/stress/go.mod b/_examples/stress/go.mod index 69be4d4..338ebf5 100644 --- a/_examples/stress/go.mod +++ b/_examples/stress/go.mod @@ -2,7 +2,10 @@ go 1.17 -require github.com/vbauerster/mpb/v8 v8.1.2 +require ( + github.com/pkg/profile v1.6.0 + github.com/vbauerster/mpb/v8 v8.1.2 +) require ( github.com/VividCortex/ewma v1.2.0 // indirect diff --git a/_examples/stress/main.go b/_examples/stress/main.go index 5815ff4..ab354b7 100644 --- a/_examples/stress/main.go +++ b/_examples/stress/main.go @@ -3,13 +3,12 @@ import ( "flag" "fmt" - "log" "math/rand" "os" - "runtime/pprof" "sync" "time" + "github.com/pkg/profile" "github.com/vbauerster/mpb/v8" "github.com/vbauerster/mpb/v8/decor" ) @@ -18,17 +17,15 @@ totalBars = 32 ) -var cpuprofile = flag.String("cpuprofile", "", "write cpu profile to file") +var proftype = flag.String("prof", "", "profile type (cpu, mem)") func main() { flag.Parse() - if *cpuprofile != "" { - f, err := os.Create(*cpuprofile) - if err != nil { - log.Fatal(err) - } - pprof.StartCPUProfile(f) - defer pprof.StopCPUProfile() + switch *proftype { + case "cpu": + defer profile.Start(profile.CPUProfile, profile.ProfilePath("."), profile.NoShutdownHook).Stop() + case "mem": + defer profile.Start(profile.MemProfile, profile.ProfilePath("."), profile.NoShutdownHook).Stop() } var wg sync.WaitGroup // passed wg will be accounted at p.Wait() call