cmd/phpgrep: add -memprofile and -cpuprofile flags

Fixes #15

Signed-off-by: Iskander Sharipov <quasilyte@gmail.com>

cmd/phpgrep/main.go

@@ -18,6 +18,9 @@ type arguments struct {
 	multiline bool
 	abs       bool
 
+	cpuProfile string
+	memProfile string
+
 	target  string
 	pattern string
 	filters []string
@@ -40,9 +43,11 @@ func main() {
 		fn   func() error
 	}{
 		{"validate flags", p.validateFlags},
+		{"start profiling", p.startProfiling},
 		{"compile filters", p.compileFilters},
 		{"compile pattern", p.compilePattern},
 		{"execute pattern", p.executePattern},
+		{"finish profiling", p.finishProfiling},
 	}
 
 	for _, step := range steps {
@@ -89,7 +94,7 @@ For more info and examples visit https://github.com/quasilyte/phpgrep
 
 Supported command-line flags:
 `
-		fmt.Fprintf(flag.CommandLine.Output(), usage)
+		fmt.Fprint(flag.CommandLine.Output(), usage)
 		flag.PrintDefaults()
 	}
 
@@ -101,6 +106,10 @@ Supported command-line flags:
 		`print absolute filenames in the output`)
 	flag.IntVar(&args.workers, "workers", 8,
 		`set the number of concurrent workers`)
+	flag.StringVar(&args.memProfile, "memprofile", "",
+		`write memory profile to the specified file`)
+	flag.StringVar(&args.cpuProfile, "cpuprofile", "",
+		`write CPU profile to the specified file`)
 
 	flag.Parse()
 

cmd/phpgrep/program.go

@@ -1,10 +1,14 @@
 package main
 
 import (
+	"bytes"
 	"fmt"
+	"io/ioutil"
 	"log"
 	"os"
 	"path/filepath"
+	"runtime"
+	"runtime/pprof"
 	"strings"
 	"sync"
 	"sync/atomic"
@@ -24,6 +28,8 @@ type program struct {
 	workers []*worker
 	filters []phpgrep.Filter
 	matches int64
+
+	cpuProfile bytes.Buffer
 }
 
 func (p *program) validateFlags() error {
@@ -43,6 +49,42 @@ func (p *program) validateFlags() error {
 	return nil
 }
 
+func (p *program) startProfiling() error {
+	if p.args.cpuProfile == "" {
+		return nil
+	}
+
+	if err := pprof.StartCPUProfile(&p.cpuProfile); err != nil {
+		return fmt.Errorf("could not start CPU profile: %v", err)
+	}
+
+	return nil
+}
+
+func (p *program) finishProfiling() error {
+	if p.args.cpuProfile != "" {
+		pprof.StopCPUProfile()
+		err := ioutil.WriteFile(p.args.cpuProfile, p.cpuProfile.Bytes(), 0666)
+		if err != nil {
+			return fmt.Errorf("write CPU profile: %v", err)
+		}
+	}
+
+	if p.args.memProfile != "" {
+		f, err := os.Create(p.args.memProfile)
+		if err != nil {
+			return fmt.Errorf("create mem profile: %v", err)
+		}
+		defer f.Close()
+		runtime.GC() // get up-to-date statistics
+		if err := pprof.WriteHeapProfile(f); err != nil {
+			return fmt.Errorf("write mem profile: %v", err)
+		}
+	}
+
+	return nil
+}
+
 func (p *program) compileFilters() error {
 	for _, s := range p.args.filters {
 		f, err := compileFilter(s)