Add profile fractions for Mutex and Block profiling (#529)

* Add profile fractions for Mutex and Block profiling

These two options need to be set in order to use the mutext contention and blocking profiling modes of pprof.

* Coverage Tax 01

* Check added flags

Author: Paulo Janotti

internal/pprofserver/pprofserver.go

@@ -18,14 +18,17 @@ import (
 	"flag"
 	"net/http"
 	_ "net/http/pprof" // Needed to enable the performance profiler
+	"runtime"
 	"strconv"
 
 	"github.com/spf13/viper"
 	"go.uber.org/zap"
 )
 
 const (
-	httpPprofPortCfg = "http-pprof-port"
+	httpPprofPortCfg          = "http-pprof-port"
+	pprofBlockProfileFraction = "pprof-block-profile-fraction"
+	pprofMutexProfileFraction = "pprof-mutex-profile-fraction"
 )
 
 // AddFlags add the command-line flags used to control the Performance Profiler
@@ -35,6 +38,14 @@ func AddFlags(flags *flag.FlagSet) {
 		httpPprofPortCfg,
 		0,
 		"Port to be used by golang net/http/pprof (Performance Profiler), the profiler is disabled if no port or 0 is specified.")
+	flags.Int(
+		pprofBlockProfileFraction,
+		-1,
+		"Fraction of blocking events that are profiled. A value <= 0 disables profiling. See runtime.SetBlockProfileRate for details.")
+	flags.Int(
+		pprofMutexProfileFraction,
+		-1,
+		"Fraction of mutex contention events that are profiled. A value <= 0 disables profiling. See runtime.SetMutexProfileFraction for details.")
 }
 
 // SetupFromViper sets up the Performance Profiler (pprof) as an HTTP endpoint
@@ -45,6 +56,9 @@ func SetupFromViper(asyncErrorChannel chan<- error, v *viper.Viper, logger *zap.
 		return nil
 	}
 
+	runtime.SetBlockProfileRate(v.GetInt(pprofBlockProfileFraction))
+	runtime.SetMutexProfileFraction(v.GetInt(pprofMutexProfileFraction))
+
 	logger.Info("Starting net/http/pprof server", zap.Int("port", port))
 	go func() {
 		if err := http.ListenAndServe(":"+strconv.Itoa(port), nil); err != http.ErrServerClosed {

internal/pprofserver/pprofserver_test.go

@@ -14,4 +14,61 @@
 
 package pprofserver
 
-// TODO: Add tests
+import (
+	"flag"
+	"net/http"
+	"runtime"
+	"strconv"
+	"testing"
+	"time"
+
+	"github.com/spf13/viper"
+	"go.uber.org/zap"
+)
+
+func TestPerformanceProfilerFlags(t *testing.T) {
+	fs := flag.NewFlagSet("test", flag.ExitOnError)
+	AddFlags(fs)
+
+	args := []string{
+		"--" + httpPprofPortCfg + "=1777",
+		"--" + pprofBlockProfileFraction + "=5",
+		"--" + pprofMutexProfileFraction + "=-1",
+	}
+
+	if err := fs.Parse(args); err != nil {
+		t.Fatalf("failed to parse arguments: %v", err)
+	}
+}
+
+func TestPerformanceProfilerServer(t *testing.T) {
+	v := viper.New()
+	const pprofPort = 17788
+	v.Set(httpPprofPortCfg, pprofPort)
+	v.Set(pprofBlockProfileFraction, 3)
+	v.Set(pprofMutexProfileFraction, 5)
+
+	asyncErrChan := make(chan error, 1)
+	if err := SetupFromViper(asyncErrChan, v, zap.NewNop()); err != nil {
+		t.Fatalf("failed to setup pprof server: %v", err)
+	}
+
+	// Give a chance for the server goroutine to run.
+	runtime.Gosched()
+
+	client := &http.Client{}
+	resp, err := client.Get("http://localhost:" + strconv.Itoa(pprofPort) + "/debug/pprof")
+	if err != nil {
+		t.Fatalf("failed to get a response from pprof server: %v", err)
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode != http.StatusOK {
+		t.Fatalf("ppropf server response: got %v want %v", resp.StatusCode, http.StatusOK)
+	}
+
+	select {
+	case err := <-asyncErrChan:
+		t.Fatalf("async err received from pprof: %v", err)
+	case <-time.After(250 * time.Millisecond):
+	}
+}