diff --git a/tools/vcli/bw/repl/repl.go b/tools/vcli/bw/repl/repl.go index c997522f..d3542f03 100644 --- a/tools/vcli/bw/repl/repl.go +++ b/tools/vcli/bw/repl/repl.go @@ -22,6 +22,9 @@ import ( "fmt" "io" "os" + "runtime" + "runtime/pprof" + "strconv" "strings" "time" @@ -88,10 +91,47 @@ func SimpleReadLine(done chan bool) <-chan string { return c } +// startProfiling tries to start pprof profiling, returning the two files to which the profiling metrics will be printed. +func startProfiling() (*os.File, *os.File, error) { + cpuProfile, err := os.Create("cpuprofile") + if err != nil { + return nil, nil, err + } + memProfile, err := os.Create("memprofile") + if err != nil { + return nil, nil, err + } + + err = pprof.StartCPUProfile(cpuProfile) + if err != nil { + return nil, nil, err + } + + return cpuProfile, memProfile, nil +} + +// stopProfiling stops pprof profiling, writing the final profiling metrics to the correspondent files and closing them. +func stopProfiling(cpuProfile, memProfile *os.File) { + if cpuProfile == nil || memProfile == nil { + fmt.Println("cpuProfile and memProfile must be both non-nil to stop profiling.") + return + } + + pprof.StopCPUProfile() + err := pprof.WriteHeapProfile(memProfile) + if err != nil { + fmt.Println(err) + } + + cpuProfile.Close() + memProfile.Close() +} + // REPL starts a read-evaluation-print-loop to run BQL commands. func REPL(od storage.Store, input *os.File, rl ReadLiner, chanSize, bulkSize, builderSize int, done chan bool) int { var tracer io.Writer - ctx, isTracingToFile, sessionStart := context.Background(), false, time.Now() + ctx, isTracingToFile, isProfiling, sessionStart := context.Background(), false, false, time.Now() + var cpuProfile, memProfile *os.File driverPlain := func() storage.Store { return od @@ -125,6 +165,11 @@ func REPL(od storage.Store, input *os.File, rl ReadLiner, chanSize, bulkSize, bu for l := range rl(done) { if strings.HasPrefix(l, "quit") { + if isProfiling { + fmt.Println("Stopping profiling and closing correspondent files.") + stopProfiling(cpuProfile, memProfile) + isProfiling = false + } done <- true break } @@ -175,6 +220,61 @@ func REPL(od storage.Store, input *os.File, rl ReadLiner, chanSize, bulkSize, bu done <- false continue } + if strings.HasPrefix(l, "start profiling") { + if isProfiling { + fmt.Println("Profiling is already ongoing.") + done <- false + continue + } + args := strings.Split(strings.TrimSpace(l)[:len(l)-1], " ") + var err error + switch len(args) { + case 2: + cpuProfile, memProfile, err = startProfiling() + if err != nil { + fmt.Println(err) + fmt.Println("Profiling failed to start.") + break + } + isProfiling = true + fmt.Println("Profiling with pprof is on.") + case 4: + if args[2] != "-cpurate" { + fmt.Printf("Invalid syntax with %q.\n\tstart profiling -cpurate \n", args[2]) + break + } + cpuProfRate, err := strconv.ParseInt(args[3], 10, 32) + if err != nil { + fmt.Println(err) + fmt.Println("Profiling failed to start.") + break + } + runtime.SetCPUProfileRate(int(cpuProfRate)) + cpuProfile, memProfile, err = startProfiling() + if err != nil { + fmt.Println(err) + fmt.Println("Profiling failed to start.") + break + } + isProfiling = true + fmt.Printf("Profiling with pprof is on (CPU profiling rate: %d samples per second).\n", cpuProfRate) + default: + fmt.Println("Invalid syntax.\n\tstart profiling -cpurate ") + } + done <- false + continue + } + if strings.HasPrefix(l, "stop profiling") { + if isProfiling { + stopProfiling(cpuProfile, memProfile) + isProfiling = false + fmt.Println("Profiling with pprof is turned off.") + } else { + fmt.Println("Profiling with pprof is already off.") + } + done <- false + continue + } if strings.HasPrefix(l, "export") { now := time.Now() args := strings.Split("bw "+strings.TrimSpace(l)[:len(l)-1], " ") @@ -255,6 +355,9 @@ func printHelp() { fmt.Println("run - runs all the BQL statements in the file.") fmt.Println("start tracing [trace_file] - starts tracing queries.") fmt.Println("stop tracing - stops tracing queries.") + fmt.Println("start profiling - starts pprof profiling for queries.") + fmt.Println("start profiling -cpurate - starts pprof profiling with customized CPU sampling rate.") + fmt.Println("stop profiling - stops pprof profiling for queries.") fmt.Println("quit - quits the console.") fmt.Println() }