better tests

Signed-off-by: Tim Vaillancourt <[email protected]>

doc/interactive-commands.md

@@ -17,7 +17,7 @@ Both interfaces may serve at the same time. Both respond to simple text command,
 - `help`: shows a brief list of available commands
 - `status`: returns a detailed status summary of migration progress and configuration
 - `sup`: returns a brief status summary of migration progress
-- `cpu-profile`: returns a base64-encoded `runtime/pprof` CPU profile with optional duration, default `30s`
+- `cpu-profile`: returns a base64-encoded runtime/pprof CPU profile with options, default '30s'. Comma-separated options 'gzip' and/or 'blocking' (blocking profile) may follow a profile duration
 - `coordinates`: returns recent (though not exactly up to date) binary log coordinates of the inspected server
 - `applier`: returns the hostname of the applier
 - `inspector`: returns the hostname of the inspector

go/logic/server.go

@@ -8,12 +8,14 @@ package logic
 import (
 	"bufio"
 	"bytes"
+	"compress/gzip"
 	"encoding/base64"
 	"errors"
 	"fmt"
 	"io"
 	"net"
 	"os"
+	"runtime"
 	"runtime/pprof"
 	"strconv"
 	"strings"
@@ -25,7 +27,10 @@ import (
 
 const defaultCPUProfileDuration = time.Second * 30
 
-var ErrCPUProfilingInProgress = errors.New("cpu profiling already in progress")
+var (
+	ErrCPUProfilingBadOption  = errors.New("unrecognized cpu profiling option")
+	ErrCPUProfilingInProgress = errors.New("cpu profiling already in progress")
+)
 
 type printStatusFunc func(PrintStatusRule, io.Writer)
 
@@ -47,24 +52,42 @@ func NewServer(migrationContext *base.MigrationContext, hooksExecutor *HooksExec
 	}
 }
 
-func (this *Server) runCPUProfile(arg string) (string, error) {
+func (this *Server) runCPUProfile(args string) (string, error) {
 	if atomic.LoadInt64(&this.isCPUProfiling) > 0 {
 		return "", ErrCPUProfilingInProgress
 	}
 
-	var err error
 	duration := defaultCPUProfileDuration
-	if arg != "" {
-		if duration, err = time.ParseDuration(arg); err != nil {
+
+	var err error
+	var useGzip bool
+	if args != "" {
+		s := strings.Split(args, ",")
+		if duration, err = time.ParseDuration(s[0]); err != nil {
 			return "", err
 		}
+		for _, arg := range s[1:] {
+			switch arg {
+			case "gzip":
+				useGzip = true
+			case "block", "blocked", "blocking":
+				runtime.SetBlockProfileRate(1)
+			default:
+				return "", ErrCPUProfilingBadOption
+			}
+		}
 	}
 
 	atomic.StoreInt64(&this.isCPUProfiling, 1)
 	defer atomic.StoreInt64(&this.isCPUProfiling, 0)
 
 	var buf bytes.Buffer
-	if err = pprof.StartCPUProfile(&buf); err != nil {
+	var writer io.Writer
+	writer = &buf
+	if useGzip {
+		writer = gzip.NewWriter(&buf)
+	}
+	if err = pprof.StartCPUProfile(writer); err != nil {
 		return "", err
 	}
 
@@ -182,7 +205,7 @@ func (this *Server) applyServerCommand(command string, writer *bufio.Writer) (pr
 			fmt.Fprint(writer, `available commands:
 status                               # Print a detailed status message
 sup                                  # Print a short status message
-cpu-profile=<duration>               # Print a base64-encoded runtime/pprof CPU profile with optional duration, default 30s
+cpu-profile=<options>                # Print a base64-encoded runtime/pprof CPU profile with options, default '30s'. Comma-separated options 'gzip' and/or 'blocking' (blocking profile) may follow a profile duration
 coordinates                          # Print the currently inspected coordinates
 applier                              # Print the hostname of the applier
 inspector                            # Print the hostname of the inspector

go/logic/server_test.go

@@ -8,23 +8,34 @@ import (
 	"github.com/openark/golib/tests"
 )
 
-func TestRunCPUProfile(t *testing.T) {
-	// failed: bad arg
-	{
+func TestServerRunCPUProfile(t *testing.T) {
+	t.Parallel()
+
+	t.Run("failed bad arg", func(t *testing.T) {
 		s := &Server{isCPUProfiling: 0}
 		profile, err := s.runCPUProfile("should-fail")
 		tests.S(t).ExpectNotNil(err)
 		tests.S(t).ExpectEquals(profile, "")
-	}
-	// failed: already running
-	{
+	})
+
+	t.Run("failed already running", func(t *testing.T) {
 		s := &Server{isCPUProfiling: 1}
 		profile, err := s.runCPUProfile("15ms")
 		tests.S(t).ExpectEquals(err, ErrCPUProfilingInProgress)
 		tests.S(t).ExpectEquals(profile, "")
-	}
-	// success
-	{
+	})
+
+	t.Run("failed with bad option", func(t *testing.T) {
+		s := &Server{
+			isCPUProfiling:   0,
+			migrationContext: base.NewMigrationContext(),
+		}
+		profile, err := s.runCPUProfile("10ms,badoption")
+		tests.S(t).ExpectEquals(err, ErrCPUProfilingBadOption)
+		tests.S(t).ExpectEquals(profile, "")
+	})
+
+	t.Run("success", func(t *testing.T) {
 		s := &Server{
 			isCPUProfiling:   0,
 			migrationContext: base.NewMigrationContext(),
@@ -37,5 +48,35 @@ func TestRunCPUProfile(t *testing.T) {
 		tests.S(t).ExpectNil(err)
 		tests.S(t).ExpectNotEquals(len(data), 0)
 		tests.S(t).ExpectEquals(s.isCPUProfiling, int64(0))
-	}
+	})
+
+	t.Run("success with block", func(t *testing.T) {
+		s := &Server{
+			isCPUProfiling:   0,
+			migrationContext: base.NewMigrationContext(),
+		}
+		profile, err := s.runCPUProfile("10ms,block")
+		tests.S(t).ExpectNil(err)
+		tests.S(t).ExpectNotEquals(profile, "")
+
+		data, err := base64.StdEncoding.DecodeString(profile)
+		tests.S(t).ExpectNil(err)
+		tests.S(t).ExpectNotEquals(len(data), 0)
+		tests.S(t).ExpectEquals(s.isCPUProfiling, int64(0))
+	})
+
+	t.Run("success with block and gzip", func(t *testing.T) {
+		s := &Server{
+			isCPUProfiling:   0,
+			migrationContext: base.NewMigrationContext(),
+		}
+		profile, err := s.runCPUProfile("10ms,block,gzip")
+		tests.S(t).ExpectNil(err)
+		tests.S(t).ExpectNotEquals(profile, "")
+
+		data, err := base64.StdEncoding.DecodeString(profile)
+		tests.S(t).ExpectNil(err)
+		tests.S(t).ExpectNotEquals(len(data), 0)
+		tests.S(t).ExpectEquals(s.isCPUProfiling, int64(0))
+	})
 }