Merge pull request #2936 from influxdb/fix-2935

fix #2935: hook CPU and memory profiling back up

CHANGELOG.md

@@ -7,6 +7,7 @@
 - [#2943](https://github.com/influxdb/influxdb/issues/2943): Ensure default retention policies are fully replicated
 - [#2948](https://github.com/influxdb/influxdb/issues/2948): Field mismatch error message to include measurement name
 - [#2919](https://github.com/influxdb/influxdb/issues/2919): Unable to insert negative floats
+- [#2935](https://github.com/influxdb/influxdb/issues/2935): Hook CPU and memory profiling back up.
 
 ## v0.9.0 [2015-06-11]
 

cmd/influxd/main.go

@@ -4,12 +4,8 @@ import (
 	"flag"
 	"fmt"
 	"io"
-	"log"
 	"math/rand"
 	"os"
-	"os/signal"
-	"runtime"
-	"runtime/pprof"
 	"strings"
 	"time"
 
@@ -135,53 +131,6 @@ func ParseCommandName(args []string) (string, []string) {
 	return "", args
 }
 
-// prof stores the file locations of active profiles.
-var prof struct {
-	cpu *os.File
-	mem *os.File
-}
-
-// StartProfile initializes the cpu and memory profile, if specified.
-func StartProfile(cpuprofile, memprofile string) {
-	if cpuprofile != "" {
-		f, err := os.Create(cpuprofile)
-		if err != nil {
-			log.Fatalf("cpuprofile: %v", err)
-		}
-		prof.cpu = f
-		pprof.StartCPUProfile(prof.cpu)
-	}
-
-	if memprofile != "" {
-		f, err := os.Create(memprofile)
-		if err != nil {
-			log.Fatalf("memprofile: %v", err)
-		}
-		prof.mem = f
-		runtime.MemProfileRate = 4096
-	}
-
-	go func() {
-		c := make(chan os.Signal, 1)
-		signal.Notify(c, os.Interrupt)
-		<-c
-		StopProfile()
-		os.Exit(0)
-	}()
-}
-
-// StopProfile closes the cpu and memory profiles if they are running.
-func StopProfile() {
-	if prof.cpu != nil {
-		pprof.StopCPUProfile()
-		prof.cpu.Close()
-	}
-	if prof.mem != nil {
-		pprof.Lookup("heap").WriteTo(prof.mem, 0)
-		prof.mem.Close()
-	}
-}
-
 // Command represents the command executed by "influxd version".
 type VersionCommand struct {
 	Stdout io.Writer

cmd/influxd/run/command.go

@@ -93,6 +93,8 @@ func (cmd *Command) Run(args ...string) error {
 	}
 	s.Commit = cmd.Commit
 	s.Version = cmd.Version
+	s.CPUProfile = options.CPUProfile
+	s.MemProfile = options.MemProfile
 	if err := s.Open(); err != nil {
 		return fmt.Errorf("open server: %s", err)
 	}

cmd/influxd/run/server.go

@@ -6,7 +6,10 @@ import (
 	"log"
 	"net"
 	"net/http"
+	"os"
+	"os/signal"
 	"runtime"
+	"runtime/pprof"
 	"time"
 
 	"github.com/influxdb/influxdb/cluster"
@@ -55,6 +58,10 @@ type Server struct {
 
 	// Server reporting
 	reportingDisabled bool
+
+	// Profiling
+	CPUProfile string
+	MemProfile string
 }
 
 // NewServer returns a new instance of Server built from a config.
@@ -247,6 +254,9 @@ func (s *Server) Err() <-chan error { return s.err }
 // Open opens the meta and data store and all services.
 func (s *Server) Open() error {
 	if err := func() error {
+		// Start profiling, if set.
+		startProfile(s.CPUProfile, s.MemProfile)
+
 		// Resolve host to address.
 		_, port, err := net.SplitHostPort(s.BindAddress)
 		if err != nil {
@@ -316,6 +326,8 @@ func (s *Server) Open() error {
 
 // Close shuts down the meta and data stores and all services.
 func (s *Server) Close() error {
+	stopProfile()
+
 	if s.Listener != nil {
 		s.Listener.Close()
 	}
@@ -409,3 +421,54 @@ type Service interface {
 	Open() error
 	Close() error
 }
+
+// prof stores the file locations of active profiles.
+var prof struct {
+	cpu *os.File
+	mem *os.File
+}
+
+// StartProfile initializes the cpu and memory profile, if specified.
+func startProfile(cpuprofile, memprofile string) {
+	if cpuprofile != "" {
+		f, err := os.Create(cpuprofile)
+		if err != nil {
+			log.Fatalf("cpuprofile: %v", err)
+		}
+		log.Printf("writing CPU profile to: %s\n", cpuprofile)
+		prof.cpu = f
+		pprof.StartCPUProfile(prof.cpu)
+	}
+
+	if memprofile != "" {
+		f, err := os.Create(memprofile)
+		if err != nil {
+			log.Fatalf("memprofile: %v", err)
+		}
+		log.Printf("writing mem profile to: %s\n", memprofile)
+		prof.mem = f
+		runtime.MemProfileRate = 4096
+	}
+
+	go func() {
+		c := make(chan os.Signal, 1)
+		signal.Notify(c, os.Interrupt)
+		<-c
+		stopProfile()
+		os.Exit(0)
+	}()
+}
+
+// StopProfile closes the cpu and memory profiles if they are running.
+func stopProfile() {
+	if prof.cpu != nil {
+		pprof.StopCPUProfile()
+		prof.cpu.Close()
+		log.Println("CPU profile stopped")
+	}
+	if prof.mem != nil {
+		pprof.Lookup("heap").WriteTo(prof.mem, 0)
+		prof.mem.Close()
+		log.Println("mem profile stopped")
+	}
+}