Skip to content

Commit

Permalink
util: pprof: Improve pprof for ease of use
Browse files Browse the repository at this point in the history
This makes it slightly cleaner.
  • Loading branch information
purpleidea committed Jan 31, 2025
1 parent 95f353c commit 7331d3a
Show file tree
Hide file tree
Showing 2 changed files with 16 additions and 13 deletions.
20 changes: 10 additions & 10 deletions main.go
Original file line number Diff line number Diff line change
Expand Up @@ -77,16 +77,6 @@ func main() {
return // for safety
}

// Run profiling if it's activated.
// TODO: Should we pass a logger into this?
ctx, cancel := context.WithCancel(context.Background())
if err := pprof.Run(ctx); err != nil {
fmt.Println(err)
os.Exit(1)
//return // redundant
}
defer cancel()

cliUtil.LogSetup(debug)
data := &cliUtil.Data{
Program: program,
Expand All @@ -102,6 +92,16 @@ func main() {
Args: os.Args,
}

// Run profiling if it's activated.
// TODO: Should we pass a logger into this?
ctx, cancel := context.WithCancel(context.Background())
if err := pprof.Run(ctx); err != nil {
fmt.Println(err)
os.Exit(1)
//return // redundant
}
defer cancel()

name := ""
if len(data.Args) > 1 {
name = data.Args[1]
Expand Down
9 changes: 6 additions & 3 deletions util/pprof/pprof.go
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,7 @@ import (
// var, it returns nil. If this is not able to start logging, then it errors. If
// it starts logging, this waits for an exit signal in a goroutine and returns
// nil. The magic env name is MGMT_PPROF_PATH. Example usage:
// MGMT_PPROF_PATH="~/pprof/out.pprof" ./mgmt run lang examples/lang/hello.mcl
// MGMT_PPROF_PATH=~/pprof/"out.pprof" ./mgmt run lang examples/lang/hello0.mcl
// go tool pprof -no_browser -http :10000 ~/pprof/out.pprof
func Run(ctx context.Context) error {
s := os.Getenv("MGMT_PPROF_PATH")
Expand All @@ -55,10 +55,12 @@ func Run(ctx context.Context) error {
log.Printf(format, v...) // XXX: use parent logger when available
}

if s == "" || !strings.HasPrefix(s, "/") {
if s == "" {
return nil // not activated
}
logf("pprof logging to: %s", s)
if !strings.HasPrefix(s, "/") {
return fmt.Errorf("pprof path is not absolute")
}

f, err := os.Create(s)
if err != nil {
Expand All @@ -67,6 +69,7 @@ func Run(ctx context.Context) error {
if err := pprof.StartCPUProfile(f); err != nil {
return fmt.Errorf("could not start CPU profile: %v", err)
}
logf("pprof logging to: %s", s)

signals := make(chan os.Signal, 1+1) // 1 * ^C + 1 * SIGTERM
signal.Notify(signals, os.Interrupt) // catch ^C
Expand Down

0 comments on commit 7331d3a

Please sign in to comment.