add --cpu-profile option, for writing out a cpu profile

Author: Brian Picciano

config/config.go

@@ -17,6 +17,7 @@ var (
 	RedisPoolSize      int
 	Debug              bool
 	BGPushPoolSize     int
+	CPUProfile         string
 )
 
 func init() {
@@ -61,6 +62,10 @@ func init() {
 		Description: "Number of goroutines to have processing NOBLOCK Q*PUSH commands",
 		Default:     "128",
 	})
+	l.Add(lever.Param{
+		Name:        "--cpu-profile",
+		Description: "Name of a file to write a cpu profile out to. If set the cpu profile will be written until okq is closed",
+	})
 	l.Parse()
 
 	ListenAddr, _ = l.ParamStr("--listen-addr")
@@ -71,4 +76,5 @@ func init() {
 	RedisPoolSize, _ = l.ParamInt("--redis-pool-size")
 	Debug = l.ParamFlag("--debug")
 	BGPushPoolSize, _ = l.ParamInt("--bg-push-pool-size")
+	CPUProfile, _ = l.ParamStr("--cpu-profile")
 }

server.go

@@ -4,6 +4,9 @@ import (
 	"errors"
 	"math/rand"
 	"net"
+	"os"
+	"os/signal"
+	"runtime/pprof"
 	"time"
 
 	"github.com/mc0/okq/clients"
@@ -27,6 +30,10 @@ func main() {
 
 	incomingConns := make(chan net.Conn)
 
+	if config.CPUProfile != "" {
+		go cpuProfile(config.CPUProfile)
+	}
+
 	go acceptConns(server, incomingConns)
 
 	for {
@@ -93,3 +100,19 @@ outer:
 		commands.Dispatch(client, command, args)
 	}
 }
+
+func cpuProfile(filename string) {
+	log.L.Printf("starting cpu profile, writing to %q", filename)
+	f, err := os.Create(filename)
+	if err != nil {
+		log.L.Fatalf("could not open cpu profile %q: %s", filename, err)
+	}
+	pprof.StartCPUProfile(f)
+
+	ch := make(chan os.Signal, 1)
+	signal.Notify(ch, os.Interrupt)
+	<-ch
+	pprof.StopCPUProfile()
+	signal.Stop(ch)
+	log.L.Printf("stopping cpu profile, ctrl-c again to exit")
+}