servenv: provide a global flag for profiling

go/vt/servenv/pprof.go

@@ -18,27 +18,235 @@ package servenv
 
 import (
 	"flag"
+	"fmt"
+	"io/ioutil"
 	"os"
+	"path/filepath"
+	"runtime"
 	"runtime/pprof"
+	"runtime/trace"
+	"strconv"
+	"strings"
+	"sync/atomic"
 
 	"vitess.io/vitess/go/vt/log"
 )
 
 var (
-	cpuProfile = flag.String("cpu_profile", "", "write cpu profile to file")
+	_         = flag.String("cpu_profile", "", "deprecated: use '-pprof=cpu' instead")
+	pprofFlag = flag.String("pprof", "", "enable profiling")
 )
 
-func init() {
-	OnInit(func() {
-		if *cpuProfile != "" {
-			f, err := os.Create(*cpuProfile)
+type profmode string
+
+const (
+	profileCPU       profmode = "cpu"
+	profileMemHeap   profmode = "mem_heap"
+	profileMemAllocs profmode = "mem_allocs"
+	profileMutex     profmode = "mutex"
+	profileBlock     profmode = "block"
+	profileTrace     profmode = "trace"
+	profileThreads   profmode = "threads"
+	profileGoroutine profmode = "goroutine"
+)
+
+func (p profmode) filename() string {
+	return fmt.Sprintf("%s.pprof", string(p))
+}
+
+type profile struct {
+	mode  profmode
+	rate  int
+	path  string
+	quiet bool
+}
+
+func parseProfileFlag(pf string) (*profile, error) {
+	if pf == "" {
+		return nil, nil
+	}
+
+	var p profile
+
+	items := strings.Split(pf, ",")
+	switch items[0] {
+	case "cpu":
+		p.mode = profileCPU
+	case "mem", "mem=heap":
+		p.mode = profileMemHeap
+		p.rate = 4096
+	case "mem=allocs":
+		p.mode = profileMemAllocs
+		p.rate = 4096
+	case "mutex":
+		p.mode = profileMutex
+		p.rate = 1
+	case "block":
+		p.mode = profileBlock
+		p.rate = 1
+	case "trace":
+		p.mode = profileTrace
+	case "threads":
+		p.mode = profileThreads
+	case "goroutine":
+		p.mode = profileGoroutine
+	default:
+		return nil, fmt.Errorf("unknown profile mode: %q", items[0])
+	}
+
+	for _, kv := range items[1:] {
+		var err error
+		fields := strings.SplitN(kv, "=", 2)
+
+		switch fields[0] {
+		case "rate":
+			if len(fields) == 1 {
+				return nil, fmt.Errorf("missing value for 'rate'")
+			}
+			p.rate, err = strconv.Atoi(fields[1])
 			if err != nil {
-				log.Fatalf("Failed to create profile file: %v", err)
+				return nil, fmt.Errorf("invalid profile rate %q: %v", fields[1], err)
 			}
-			pprof.StartCPUProfile(f)
-			OnTerm(func() {
-				pprof.StopCPUProfile()
-			})
+
+		case "path":
+			if len(fields) == 1 {
+				return nil, fmt.Errorf("missing value for 'path'")
+			}
+			p.path = fields[1]
+
+		case "quiet":
+			if len(fields) == 1 {
+				p.quiet = true
+				continue
+			}
+
+			p.quiet, err = strconv.ParseBool(fields[1])
+			if err != nil {
+				return nil, fmt.Errorf("invalid quiet flag %q: %v", fields[1], err)
+			}
+		default:
+			return nil, fmt.Errorf("unknown flag: %q", fields[0])
+		}
+	}
+
+	return &p, nil
+}
+
+var profileStarted uint32
+
+// start begins the configured profiling process and returns a cleanup function
+// that must be executed before process termination to flush the profile to disk.
+// Based on the profiling code in github.com/pkg/profile
+func (prof *profile) start() func() {
+	if !atomic.CompareAndSwapUint32(&profileStarted, 0, 1) {
+		log.Fatal("profile: Start() already called")
+	}
+
+	var (
+		path string
+		err  error
+		logf = func(format string, args ...interface{}) {}
+	)
+
+	if prof.path != "" {
+		path = prof.path
+		err = os.MkdirAll(path, 0777)
+	} else {
+		path, err = ioutil.TempDir("", "profile")
+	}
+	if err != nil {
+		log.Fatalf("pprof: could not create initial output directory: %v", err)
+	}
+
+	if !prof.quiet {
+		logf = log.Infof
+	}
+
+	fn := filepath.Join(path, prof.mode.filename())
+	f, err := os.Create(fn)
+	if err != nil {
+		log.Fatalf("pprof: could not create profile %q: %v", fn, err)
+	}
+	logf("pprof: %s profiling enabled, %s", string(prof.mode), fn)
+
+	switch prof.mode {
+	case profileCPU:
+		pprof.StartCPUProfile(f)
+		return func() {
+			pprof.StopCPUProfile()
+			f.Close()
+		}
+
+	case profileMemHeap, profileMemAllocs:
+		old := runtime.MemProfileRate
+		runtime.MemProfileRate = prof.rate
+		return func() {
+			tt := "heap"
+			if prof.mode == profileMemAllocs {
+				tt = "allocs"
+			}
+			pprof.Lookup(tt).WriteTo(f, 0)
+			f.Close()
+			runtime.MemProfileRate = old
+		}
+
+	case profileMutex:
+		runtime.SetMutexProfileFraction(prof.rate)
+		return func() {
+			if mp := pprof.Lookup("mutex"); mp != nil {
+				mp.WriteTo(f, 0)
+			}
+			f.Close()
+			runtime.SetMutexProfileFraction(0)
+		}
+
+	case profileBlock:
+		runtime.SetBlockProfileRate(prof.rate)
+		return func() {
+			pprof.Lookup("block").WriteTo(f, 0)
+			f.Close()
+			runtime.SetBlockProfileRate(0)
+		}
+
+	case profileThreads:
+		return func() {
+			if mp := pprof.Lookup("threadcreate"); mp != nil {
+				mp.WriteTo(f, 0)
+			}
+			f.Close()
+		}
+
+	case profileTrace:
+		if err := trace.Start(f); err != nil {
+			log.Fatalf("pprof: could not start trace: %v", err)
+		}
+		return func() {
+			trace.Stop()
+			f.Close()
+		}
+
+	case profileGoroutine:
+		return func() {
+			if mp := pprof.Lookup("goroutine"); mp != nil {
+				mp.WriteTo(f, 0)
+			}
+			f.Close()
+		}
+
+	default:
+		panic("unsupported profile mode")
+	}
+}
+
+func init() {
+	OnInit(func() {
+		prof, err := parseProfileFlag(*pprofFlag)
+		if err != nil {
+			log.Fatal(err)
+		}
+		if prof != nil {
+			stop := prof.start()
+			OnTerm(stop)
 		}
 	})
 }

go/vt/servenv/pprof_test.go

@@ -0,0 +1,45 @@
+package servenv
+
+import (
+	"reflect"
+	"testing"
+)
+
+func TestParseProfileFlag(t *testing.T) {
+	tests := []struct {
+		arg     string
+		want    *profile
+		wantErr bool
+	}{
+		{"", nil, false},
+		{"mem", &profile{mode: profileMemHeap, rate: 4096}, false},
+		{"mem,rate=1234", &profile{mode: profileMemHeap, rate: 1234}, false},
+		{"mem,rate", nil, true},
+		{"mem,rate=foobar", nil, true},
+		{"mem=allocs", &profile{mode: profileMemAllocs, rate: 4096}, false},
+		{"mem=allocs,rate=420", &profile{mode: profileMemAllocs, rate: 420}, false},
+		{"block", &profile{mode: profileBlock, rate: 1}, false},
+		{"block,rate=4", &profile{mode: profileBlock, rate: 4}, false},
+		{"cpu", &profile{mode: profileCPU}, false},
+		{"cpu,quiet", &profile{mode: profileCPU, quiet: true}, false},
+		{"cpu,quiet=true", &profile{mode: profileCPU, quiet: true}, false},
+		{"cpu,quiet=false", &profile{mode: profileCPU, quiet: false}, false},
+		{"cpu,quiet=foobar", nil, true},
+		{"cpu,path=", &profile{mode: profileCPU, path: ""}, false},
+		{"cpu,path", nil, true},
+		{"cpu,path=a", &profile{mode: profileCPU, path: "a"}, false},
+		{"cpu,path=a/b/c/d", &profile{mode: profileCPU, path: "a/b/c/d"}, false},
+	}
+	for _, tt := range tests {
+		t.Run(tt.arg, func(t *testing.T) {
+			got, err := parseProfileFlag(tt.arg)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("parseProfileFlag() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			if !reflect.DeepEqual(got, tt.want) {
+				t.Errorf("parseProfileFlag() got = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}

go/vt/servenv/servenv.go

@@ -33,7 +33,6 @@ import (
 	"net/url"
 	"os"
 	"os/signal"
-	"runtime"
 	"strings"
 	"sync"
 	"syscall"
@@ -56,11 +55,11 @@ var (
 	Port *int
 
 	// Flags to alter the behavior of the library.
-	lameduckPeriod       = flag.Duration("lameduck-period", 50*time.Millisecond, "keep running at least this long after SIGTERM before stopping")
-	onTermTimeout        = flag.Duration("onterm_timeout", 10*time.Second, "wait no more than this for OnTermSync handlers before stopping")
-	memProfileRate       = flag.Int("mem-profile-rate", 512*1024, "profile every n bytes allocated")
-	mutexProfileFraction = flag.Int("mutex-profile-fraction", 0, "profile every n mutex contention events (see runtime.SetMutexProfileFraction)")
-	catchSigpipe         = flag.Bool("catch-sigpipe", false, "catch and ignore SIGPIPE on stdout and stderr if specified")
+	lameduckPeriod = flag.Duration("lameduck-period", 50*time.Millisecond, "keep running at least this long after SIGTERM before stopping")
+	onTermTimeout  = flag.Duration("onterm_timeout", 10*time.Second, "wait no more than this for OnTermSync handlers before stopping")
+	_              = flag.Int("mem-profile-rate", 512*1024, "deprecated: use '-pprof=mem' instead")
+	_              = flag.Int("mutex-profile-fraction", 0, "deprecated: use '-pprof=mutex' instead")
+	catchSigpipe   = flag.Bool("catch-sigpipe", false, "catch and ignore SIGPIPE on stdout and stderr if specified")
 
 	// mutex used to protect the Init function
 	mu sync.Mutex
@@ -106,13 +105,6 @@ func Init() {
 		log.Exitf("servenv.Init: running this as root makes no sense")
 	}
 
-	runtime.MemProfileRate = *memProfileRate
-
-	if *mutexProfileFraction != 0 {
-		log.Infof("setting mutex profile fraction to %v", *mutexProfileFraction)
-		runtime.SetMutexProfileFraction(*mutexProfileFraction)
-	}
-
 	// We used to set this limit directly, but you pretty much have to
 	// use a root account to allow increasing a limit reliably. Dropping
 	// privileges is also tricky. The best strategy is to make a shell