Add server command

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

doc/interactive-commands.md

@@ -17,6 +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`
 - `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

@@ -7,17 +7,24 @@ package logic
 
 import (
 	"bufio"
+	"bytes"
+	"encoding/base64"
+	"errors"
 	"fmt"
 	"io"
 	"net"
 	"os"
+	"runtime/pprof"
 	"strconv"
 	"strings"
 	"sync/atomic"
+	"time"
 
 	"github.com/github/gh-ost/go/base"
 )
 
+var ErrCPUProfilingInProgress = errors.New("cpu profiling already in progress")
+
 type printStatusFunc func(PrintStatusRule, io.Writer)
 
 // Server listens for requests on a socket file or via TCP
@@ -27,6 +34,7 @@ type Server struct {
 	tcpListener      net.Listener
 	hooksExecutor    *HooksExecutor
 	printStatus      printStatusFunc
+	isCPUProfiling   int64
 }
 
 func NewServer(migrationContext *base.MigrationContext, hooksExecutor *HooksExecutor, printStatus printStatusFunc) *Server {
@@ -37,6 +45,25 @@ func NewServer(migrationContext *base.MigrationContext, hooksExecutor *HooksExec
 	}
 }
 
+func (this *Server) runCPUProfile(duration time.Duration) (string, error) {
+	if atomic.LoadInt64(&this.isCPUProfiling) > 0 {
+		return "", ErrCPUProfilingInProgress
+	}
+
+	atomic.StoreInt64(&this.isCPUProfiling, 1)
+	defer atomic.StoreInt64(&this.isCPUProfiling, 0)
+
+	var buf bytes.Buffer
+	if err := pprof.StartCPUProfile(&buf); err != nil {
+		return "", err
+	}
+
+	time.Sleep(duration)
+	pprof.StopCPUProfile()
+
+	return base64.StdEncoding.EncodeToString(buf.Bytes()), nil
+}
+
 func (this *Server) BindSocketFile() (err error) {
 	if this.migrationContext.ServeSocketFile == "" {
 		return nil
@@ -144,6 +171,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
 coordinates                          # Print the currently inspected coordinates
 applier                              # Print the hostname of the applier
 inspector                            # Print the hostname of the inspector
@@ -169,6 +197,18 @@ help                                 # This message
 		return ForcePrintStatusOnlyRule, nil
 	case "info", "status":
 		return ForcePrintStatusAndHintRule, nil
+	case "cpu-profile":
+		profileDuration := time.Second * 30
+		if arg != "" {
+			if profileDuration, err = time.ParseDuration(arg); err != nil {
+				return NoPrintStatusRule, err
+			}
+		}
+		profile, err := this.runCPUProfile(profileDuration)
+		if err == nil {
+			fmt.Fprintln(writer, profile)
+		}
+		return NoPrintStatusRule, err
 	case "coordinates":
 		{
 			if argIsQuestion || arg == "" {

go/logic/server_test.go

@@ -0,0 +1,27 @@
+package logic
+
+import (
+	"encoding/base64"
+	"testing"
+	"time"
+
+	"github.com/openark/golib/tests"
+)
+
+func TestRunCPUProfile(t *testing.T) {
+	{
+		s := &Server{isCPUProfiling: 0}
+		profile, err := s.runCPUProfile(time.Millisecond * 10)
+		tests.S(t).ExpectNil(err)
+		tests.S(t).ExpectNotEquals(profile, "")
+
+		_, err = base64.StdEncoding.DecodeString(profile)
+		tests.S(t).ExpectNil(err)
+	}
+	{
+		s := &Server{isCPUProfiling: 1}
+		profile, err := s.runCPUProfile(time.Millisecond * 10)
+		tests.S(t).ExpectNotNil(err)
+		tests.S(t).ExpectEquals(profile, "")
+	}
+}